fixed shape for actnorm constructor

vislearn · Aug 18, 2023 · 120c520 · 120c520
1 parent f5690e3
commit 120c520
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/FrEIA/modules/invertible_resnet.py b/FrEIA/modules/invertible_resnet.py
@@ -30,8 +30,9 @@ def __init__(self, dims_in, dims_c=None, init_data: torch.Tensor = None):
 
         self.register_buffer("is_initialized", torch.tensor(False))
 
-        self.log_scale = nn.Parameter(torch.empty(1, *dims_in))
-        self.loc = nn.Parameter(torch.empty(1, *dims_in))
+        dims = next(iter(dims_in))
+        self.log_scale = nn.Parameter(torch.empty(1, *dims))
+        self.loc = nn.Parameter(torch.empty(1, *dims))
 
         if init_data is not None:
             self.initialize(init_data)