Update correlation.py

fix CUDA_ERROR_ILLEGAL_ADDRESS

Author: 1374839016

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):