Update correlation.py

1374839016 · web-flow · commit 07c3f5a1675a · 2023-06-24T20:03:14.000+08:00
fix CUDA_ERROR_ILLEGAL_ADDRESS
diff --git a/correlation/correlation.py b/correlation/correlation.py
@@ -381,7 +381,8 @@ def backward(self, gradOutput):
 # end
 
 def FunctionCorrelation(tenOne, tenTwo):
-    return _FunctionCorrelation.apply(tenOne, tenTwo)
+    with cupy.cuda.Device(tenOne.get_device()):
+        return _FunctionCorrelation.apply(tenOne, tenTwo)
 # end
 
 class ModuleCorrelation(torch.nn.Module):
