Merge pull request #44 from chcodes/fix-device-contexts

Added missing device contexts
mikgroup · Apr 27, 2020 · cbbe549 · cbbe549
2 parents 1576270 + 4bffa1c
commit cbbe549
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/sigpy/alg.py b/sigpy/alg.py
@@ -338,11 +338,13 @@ def __init__(self, proxfc, proxg, A, AH, x, u,
 
         if self.gamma_primal > 0:
             xp = self.x_device.xp
-            self.tau_min = xp.amin(xp.abs(tau)).item()
+            with self.x_device:
+                self.tau_min = xp.amin(xp.abs(tau)).item()
 
         if self.gamma_dual > 0:
             xp = self.u_device.xp
-            self.sigma_min = xp.amin(xp.abs(sigma)).item()
+            with self.u_device:
+                self.sigma_min = xp.amin(xp.abs(sigma)).item()
 
         self.resid = np.infty
 

diff --git a/sigpy/linop.py b/sigpy/linop.py
@@ -646,7 +646,8 @@ def __init__(self, oshape, ishape):
         super().__init__(oshape, ishape)
 
     def _apply(self, input):
-        return input.reshape(self.oshape)
+        with backend.get_device(input):
+            return input.reshape(self.oshape)
 
     def _adjoint_linop(self):
         return Reshape(self.ishape, self.oshape)
@@ -1050,8 +1051,9 @@ def __init__(self, oshape, ishape, ishift=None, oshift=None):
         super().__init__(oshape, ishape)
 
     def _apply(self, input):
-        return util.resize(input, self.oshape,
-                           ishift=self.ishift, oshift=self.oshift)
+        with backend.get_device(input):
+            return util.resize(input, self.oshape,
+                               ishift=self.ishift, oshift=self.oshift)
 
     def _adjoint_linop(self):
         return Resize(self.ishape, self.oshape,