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