From c570f09681b2e30646457e8c7adcd846b75d5b47 Mon Sep 17 00:00:00 2001
From: gilbertlee-amd <44450918+gilbertlee-amd@users.noreply.github.com>
Date: Tue, 15 Dec 2020 17:37:31 -0700
Subject: [PATCH] [TransferBench] Fixing bug with fine-grained memory
 allocation (#311)

* Fixing bug with fine-grained memory

[ROCm/rccl commit: 41c35dad485d20c519a31e50c38ddb57445353df]
---
 projects/rccl/tools/TransferBench/TransferBench.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/projects/rccl/tools/TransferBench/TransferBench.cpp b/projects/rccl/tools/TransferBench/TransferBench.cpp
index f6d90fdb29..faed8dae44 100644
--- a/projects/rccl/tools/TransferBench/TransferBench.cpp
+++ b/projects/rccl/tools/TransferBench/TransferBench.cpp
@@ -764,6 +764,7 @@ void AllocateMemory(MemType memType, int devIndex, size_t numBytes, float** memP
   }
   else if (memType == MEM_GPU_FINE)
   {
+    HIP_CALL(hipSetDevice(devIndex));
     HIP_CALL(hipExtMallocWithFlags((void**)memPtr, numBytes, hipDeviceMallocFinegrained));
   }
   else