小优化

ztxz16 · Jan 20, 2025 · 03ce4f7 · 03ce4f7
1 parent ec1a067
commit 03ce4f7
Showing 1 changed file with 9 additions and 8 deletions.
diff --git a/src/fastllm.cpp b/src/fastllm.cpp
@@ -1167,14 +1167,15 @@ namespace fastllm {
                 } else if (device == DataDevice::CUDA) {
                     int sourceDevice = this->dataDeviceIds.size() == 0 ? 0 : this->dataDeviceIds[0];
                     int destDevice = deviceIds.size() == 0 ? 0 : deviceIds[0];
-                    FastllmCudaSetDevice(destDevice);
-                    void *newCudaData = FastllmCudaMalloc(expansionBytes);
-
-                    FastllmCudaMemcpyBetweenDevices(destDevice, newCudaData, sourceDevice, this->cudaData, expansionBytes);
-                    FastllmCudaSetDevice(sourceDevice);
-                    FastllmCudaFree(this->cudaData);
-                    this->cudaData = newCudaData;
-                    FastllmCudaSetDevice(destDevice);
+                    if (sourceDevice != destDevice) {
+                                        FastllmCudaSetDevice(destDevice);
+                                        void *newCudaData = FastllmCudaMalloc(expansionBytes);
+                                        FastllmCudaMemcpyBetweenDevices(destDevice, newCudaData, sourceDevice, this->cudaData, expansionBytes);
+                                        FastllmCudaSetDevice(sourceDevice);
+                                        FastllmCudaFree(this->cudaData);
+                                        this->cudaData = newCudaData;
+                                        FastllmCudaSetDevice(destDevice);
+                    }
                 }
             }
 #endif