fix type hint in squeeze [no ci]

Author: vpratz

bayesflow/adapters/adapter.py

@@ -820,7 +820,7 @@ def split(self, key: str, *, into: Sequence[str], indices_or_sections: int | Seq
 
         return self
 
-    def squeeze(self, keys: str | Sequence[str], *, axis: int | tuple):
+    def squeeze(self, keys: str | Sequence[str], *, axis: int | Sequence[int]):
         """Append a :py:class:`~transforms.Squeeze` transform to the adapter.
 
         Parameters

bayesflow/adapters/transforms/squeeze.py

@@ -1,5 +1,6 @@
 import numpy as np
 
+from collections.abc import Sequence
 from bayesflow.utils.serialization import serializable, serialize
 
 from .elementwise_transform import ElementwiseTransform
@@ -29,8 +30,10 @@ class Squeeze(ElementwiseTransform):
     It is recommended to precede this transform with a :class:`~bayesflow.adapters.transforms.ToArray` transform.
     """
 
-    def __init__(self, *, axis: int | tuple):
+    def __init__(self, *, axis: int | Sequence[int]):
         super().__init__()
+        if isinstance(axis, Sequence):
+            axis = tuple(axis)
         self.axis = axis
 
     def get_config(self) -> dict: