Fix event shape typehint

davidnabergoj · Oct 13, 2023 · 1241200 · 1241200
1 parent 1293606
commit 1241200
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/normalizing_flows/bijections/base.py b/normalizing_flows/bijections/base.py
@@ -91,7 +91,8 @@ def invert(bijection: Bijection) -> Bijection:
 
 class BijectiveComposition(Bijection):
     def __init__(self,
-                 event_shape: torch.Size, layers: List[Bijection],
+                 event_shape: Union[torch.Size, Tuple[int, ...]],
+                 layers: List[Bijection],
                  context_shape: Union[torch.Size, Tuple[int, ...]] = None):
         super().__init__(event_shape=event_shape, context_shape=context_shape)
         self.layers = nn.ModuleList(layers)