diff --git a/correlation/correlation.py b/correlation/correlation.py
index 5b560d4..4b7b1ba 100644
--- 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):