fix mask error (#398)

gelles-brandeis · Nov 19, 2022 · 034bf5e · 034bf5e
1 parent 5809962
commit 034bf5e
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/tapqir/models/cosmos.py b/tapqir/models/cosmos.py
@@ -395,7 +395,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/crosstalk.py b/tapqir/models/crosstalk.py
@@ -126,7 +126,7 @@ def model(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):
                 # background mean and std
                 background_mean = pyro.sample(

diff --git a/tapqir/models/hmm.py b/tapqir/models/hmm.py
@@ -138,7 +138,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(