Fix typehint syntaxt

davidnabergoj · Oct 19, 2024 · 38d29e1 · 38d29e1
1 parent 5aefcd6
commit 38d29e1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchflows/bijections/finite/autoregressive/conditioning/transforms.py b/torchflows/bijections/finite/autoregressive/conditioning/transforms.py
@@ -273,7 +273,7 @@ def __init__(self, *args, **kwargs):
 class FeedForward(TensorConditionerTransform):
     def __init__(self,
                  input_event_shape: Union[torch.Size, Tuple[int, ...]],
-                 parameter_shape: torch.Union[torch.Size, Tuple[int, ...]],
+                 parameter_shape: Union[torch.Size, Tuple[int, ...]],
                  context_shape: Union[torch.Size, Tuple[int, ...]] = None,
                  n_hidden: int = None,
                  n_layers: int = 2,