Fix off-by-one

vislearn · Sep 26, 2023 · ec8dab6 · ec8dab6
1 parent 84cfcb6
commit ec8dab6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/FrEIA/modules/invertible_resnet.py b/FrEIA/modules/invertible_resnet.py
@@ -31,7 +31,7 @@ def __init__(self, dims_in, dims_c=None, init_data: torch.Tensor = None):
         self.register_buffer("is_initialized", torch.tensor(False))
 
         dims = list(next(iter(dims_in)))
-        dims[2:] = [1] * len(dims[2:])
+        dims[1:] = [1] * len(dims[1:])
         self.log_scale = nn.Parameter(torch.empty(1, *dims))
         self.loc = nn.Parameter(torch.empty(1, *dims))