small fix for distributed

NVIDIA · Jan 22, 2025 · 33449fd · 33449fd
1 parent 0fe5678
commit 33449fd
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torch_harmonics/distributed/distributed_convolution.py b/torch_harmonics/distributed/distributed_convolution.py
@@ -213,7 +213,7 @@ def __init__(
         out_channels: int,
         in_shape: Tuple[int],
         out_shape: Tuple[int],
-        kernel_shape: Union[int, List[int]],
+        kernel_shape: Union[int, Tuple[int], Tuple[int, int]],
         basis_type: Optional[str] = "piecewise linear",
         basis_norm_mode: Optional[str] = "mean",
         groups: Optional[int] = 1,
@@ -354,7 +354,7 @@ def __init__(
         out_channels: int,
         in_shape: Tuple[int],
         out_shape: Tuple[int],
-        kernel_shape: Union[int, List[int]],
+        kernel_shape: Union[int, Tuple[int], Tuple[int, int]],
         basis_type: Optional[str] = "piecewise linear",
         basis_norm_mode: Optional[str] = "mean",
         groups: Optional[int] = 1,