unwrap input in dimension check AllInOne block

vislearn · Aug 18, 2023 · e4b1253 · e4b1253
1 parent 5896c1b
commit e4b1253
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/FrEIA/modules/all_in_one_block.py b/FrEIA/modules/all_in_one_block.py
@@ -239,7 +239,7 @@ def _affine(self, x, a, rev=False):
 
     def forward(self, x, c=[], rev=False, jac=True):
         '''See base class docstring'''
-        if x.shape[1:] != self.dims_in[0][1:]:
+        if x.shape[0][1:] != self.dims_in[0][1:]:
             raise RuntimeError(f"Expected input of shape {self.dims_in[0]}, "
                              f"got {x.shape}.")
         if self.householder: