From 07c3f5a1675ad09bb073c4a03a480709c9a74fa8 Mon Sep 17 00:00:00 2001
From: Jianxin Hu <51019236+1374839016@users.noreply.github.com>
Date: Sat, 24 Jun 2023 20:03:14 +0800
Subject: [PATCH] Update correlation.py

fix CUDA_ERROR_ILLEGAL_ADDRESS
---
 correlation/correlation.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

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