diff --git a/mooncake-transfer-engine/example/transfer_engine_bench.cpp b/mooncake-transfer-engine/example/transfer_engine_bench.cpp
index 6d763643..44ecd4a1 100644
--- a/mooncake-transfer-engine/example/transfer_engine_bench.cpp
+++ b/mooncake-transfer-engine/example/transfer_engine_bench.cpp
@@ -32,8 +32,8 @@
 
 #ifdef USE_CUDA
 #include <bits/stdint-uintn.h>
-#include <cuda_runtime.h>
 #include <cuda.h>
+#include <cuda_runtime.h>
 
 #ifdef USE_NVMEOF
 #include <cufile.h>
@@ -100,12 +100,12 @@ static void *allocateMemoryPool(size_t size, int socket_id,
         int gpu_id = FLAGS_gpu_id;
         void *d_buf;
         checkCudaError(cudaSetDevice(gpu_id), "Failed to set device");
-#ifdef USE_NVLINK
-        d_buf = mooncake::NvlinkTransport::allocatePinnedLocalMemory(size);
-#else
-        checkCudaError(cudaMalloc(&d_buf, size),
-                       "Failed to allocate device memory");
-#endif
+        if (FLAGS_protocol == "nvlink") {
+            d_buf = mooncake::NvlinkTransport::allocatePinnedLocalMemory(size);
+        } else {
+            checkCudaError(cudaMalloc(&d_buf, size),
+                           "Failed to allocate device memory");
+        }
         return d_buf;
     }
 #endif
@@ -114,14 +114,14 @@ static void *allocateMemoryPool(size_t size, int socket_id,
 
 static void freeMemoryPool(void *addr, size_t size) {
 #ifdef USE_CUDA
-#ifdef USE_NVLINK
-    CUmemGenericAllocationHandle handle;
-    auto result = cuMemRetainAllocationHandle(&handle, addr);
-    if (result == CUDA_SUCCESS) {
-        mooncake::NvlinkTransport::freePinnedLocalMemory(addr);
-        return;
+    if (FLAGS_protocol == "nvlink") {
+        CUmemGenericAllocationHandle handle;
+        auto result = cuMemRetainAllocationHandle(&handle, addr);
+        if (result == CUDA_SUCCESS) {
+            mooncake::NvlinkTransport::freePinnedLocalMemory(addr);
+            return;
+        }
     }
-#endif
     // check pointer on GPU
     cudaPointerAttributes attributes;
     checkCudaError(cudaPointerGetAttributes(&attributes, addr),
@@ -406,12 +406,7 @@ int target() {
     buffer_num = FLAGS_use_vram ? 1 : NR_SOCKETS;
     if (FLAGS_use_vram) LOG(INFO) << "VRAM is used";
     for (int i = 0; i < buffer_num; ++i) {
-#ifdef USE_NVLINK
-        addr[i] = mooncake::NvlinkTransport::allocatePinnedLocalMemory(
-            FLAGS_buffer_size);
-#else
         addr[i] = allocateMemoryPool(FLAGS_buffer_size, i, FLAGS_use_vram);
-#endif
         std::string name_prefix = FLAGS_use_vram ? "cuda:" : "cpu:";
         int rc = engine->registerLocalMemory(addr[i], FLAGS_buffer_size,
                                              name_prefix + std::to_string(i));
@@ -431,11 +426,7 @@ int target() {
     while (target_running) sleep(1);
     for (int i = 0; i < buffer_num; ++i) {
         engine->unregisterLocalMemory(addr[i]);
-#ifdef USE_NVLINK
-        mooncake::NvlinkTransport::freePinnedLocalMemory(addr[i]);
-#else
         freeMemoryPool(addr[i], FLAGS_buffer_size);
-#endif
     }
 
     return 0;
diff --git a/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp b/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp
index 7219b253..5312b680 100644
--- a/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp
+++ b/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp
@@ -281,6 +281,12 @@ int NvlinkTransport::registerLocalMemory(void *addr, size_t length,
             return 0;
         }
 
+        cudaError_t err = cudaSetDevice(0);
+        if (err != cudaSuccess) {
+            LOG(ERROR) << "NvlinkTransport: cudaSetDevice failed";
+            return -1;
+        }
+
         // Find whole physical page for memory registration
         void *real_addr;
         size_t real_size;