fixed typing in AIO input check

vislearn · Aug 18, 2023 · 31ef2b8 · 31ef2b8
1 parent b983024
commit 31ef2b8
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/FrEIA/modules/all_in_one_block.py b/FrEIA/modules/all_in_one_block.py
@@ -239,9 +239,9 @@ def _affine(self, x, a, rev=False):
 
     def forward(self, x, c=[], rev=False, jac=True):
         '''See base class docstring'''
-        if x[0].shape[1:] != self.dims_in[0][1:]:
+        if tuple(x[0].shape[1:]) != self.dims_in[0]:
             raise RuntimeError(f"Expected input of shape {self.dims_in[0]}, "
-                             f"got {x.shape}.")
+                             f"got {tuple(x[0].shape[1:])}.")
         if self.householder:
             self.w_perm = self._construct_householder_permutation()
             if rev or self.reverse_pre_permute: