fix fetch device (#399)

gelles-brandeis · Nov 20, 2022 · 900973e · 900973e
1 parent 034bf5e
commit 900973e
Show file tree

Hide file tree

Showing 5 changed files with 6 additions and 14 deletions.
diff --git a/tapqir/models/cosmos.py b/tapqir/models/cosmos.py
@@ -216,7 +216,7 @@ def model(self):
 
         with channels as cdx, aois as ndx:
             ndx = ndx[:, None, None]
-            mask = Vindex(self.data.mask)[ndx].to(self.device)
+            mask = Vindex(self.data.mask.to(self.device))[ndx]
             with handlers.mask(mask=mask):
                 # background mean and std
                 background_mean = pyro.sample(

diff --git a/tapqir/models/crosstalk.py b/tapqir/models/crosstalk.py
@@ -341,7 +341,7 @@ def guide(self):
 
         with aois as ndx:
             ndx = ndx[:, None]
-            mask = Vindex(self.data.mask)[ndx].to(self.device)
+            mask = Vindex(self.data.mask.to(self.device))[ndx]
             with handlers.mask(mask=mask):
                 pyro.sample(
                     "background_mean",

diff --git a/tapqir/models/hmm.py b/tapqir/models/hmm.py
@@ -323,7 +323,7 @@ def guide(self):
 
         with channels as cdx, aois as ndx:
             ndx = ndx[:, None, None]
-            mask = Vindex(self.data.mask)[ndx].to(self.device)
+            mask = Vindex(self.data.mask.to(self.device))[ndx]
             with handlers.mask(mask=mask):
                 pyro.sample(
                     "background_mean",

diff --git a/tapqir/models/model.py b/tapqir/models/model.py
@@ -89,14 +89,6 @@ def to(self, device: str, dtype: str = "double") -> None:
             torch.set_default_tensor_type(torch.DoubleTensor)
         else:
             torch.set_default_tensor_type(torch.FloatTensor)
-        # change loaded data device
-        if hasattr(self, "data"):
-            self.data.ontarget = self.data.ontarget._replace(device=self.device)
-            self.data.offtarget = self.data.offtarget._replace(device=self.device)
-            self.data.offset = self.data.offset._replace(
-                samples=self.data.offset.samples.to(self.device),
-                weights=self.data.offset.weights.to(self.device),
-            )
 
     @property
     def Q(self):

diff --git a/tapqir/utils/dataset.py b/tapqir/utils/dataset.py
@@ -139,9 +139,9 @@ def median(self) -> torch.Tensor:
 
     def fetch(self, ndx, fdx, cdx):
         return (
-            Vindex(self.images)[ndx, fdx, cdx].to(self.device),
-            Vindex(self.xy)[ndx, fdx, cdx].to(self.device),
-            Vindex(self.is_ontarget)[ndx].to(self.device),
+            Vindex(self.images.to(self.device))[ndx, fdx, cdx],
+            Vindex(self.xy.to(self.device))[ndx, fdx, cdx],
+            Vindex(self.is_ontarget.to(self.device))[ndx],
         )
 
     @lazy_property