keep device in clone()

Co-authored-by: Alexandru Fikl <[email protected]>
inducer · Feb 11, 2025 · 3ee28cc · 3ee28cc
1 parent 340f9dc
commit 3ee28cc
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/arraycontext/impl/cupy/__init__.py b/arraycontext/impl/cupy/__init__.py
@@ -75,6 +75,8 @@ def __init__(self, device: int | None = None) -> None:
         super().__init__()
         self._loopy_transform_cache = {}
 
+        self.device = device
+
         if device is not None:
             import cupy as cp
             cp.cuda.runtime.setDevice(device)
@@ -88,7 +90,7 @@ def _get_fake_numpy_namespace(self):
     # {{{ ArrayContext interface
 
     def clone(self):
-        return type(self)()
+        return type(self)(self.device)
 
     @overload
     def from_numpy(self, array: np.ndarray) -> Array: