diff --git a/equilib/cube2equi/cube2equi_torch/cube2equi.py b/equilib/cube2equi/cube2equi_torch/cube2equi.py
index 9afa1688..7e8bf6a1 100644
--- a/equilib/cube2equi/cube2equi_torch/cube2equi.py
+++ b/equilib/cube2equi/cube2equi_torch/cube2equi.py
@@ -95,7 +95,7 @@ def _equirect_facetype(self, h: int, w: int) -> torch.Tensor:
 
         return tp.type(torch.int32)
 
-    def create_equi_grid(self, h_out: int, w_out: int) -> Tuple(torch.Tensor):
+    def create_equi_grid(self, h_out: int, w_out: int) -> Tuple[torch.Tensor]:
         _dtype = torch.float32
         theta = torch.linspace(-math.pi, math.pi, steps=w_out, dtype=_dtype)
         phi = torch.linspace(math.pi, -math.pi, steps=h_out, dtype=_dtype) / 2
diff --git a/equilib/equi2cube/equi2cube_numpy/utils.py b/equilib/equi2cube/equi2cube_numpy/utils.py
index cdafdd2d..a9d33af8 100644
--- a/equilib/equi2cube/equi2cube_numpy/utils.py
+++ b/equilib/equi2cube/equi2cube_numpy/utils.py
@@ -53,7 +53,7 @@ def cube_h2list(cube_h: np.ndarray) -> List[np.ndarray]:
     return np.split(cube_h, 6, axis=-1)
 
 
-def cube_h2dict(cube_h: np.ndarray) -> Dict[np.ndarray]:
+def cube_h2dict(cube_h: np.ndarray) -> Dict[str, np.ndarray]:
     cube_list = cube_h2list(cube_h)
     return dict(
         [
diff --git a/equilib/equi2cube/equi2cube_torch/equi2cube.py b/equilib/equi2cube/equi2cube_torch/equi2cube.py
index c74e540c..2562e68c 100644
--- a/equilib/equi2cube/equi2cube_torch/equi2cube.py
+++ b/equilib/equi2cube/equi2cube_torch/equi2cube.py
@@ -109,7 +109,7 @@ def run(
         torch.Tensor,
         List[torch.Tensor],
         List[Dict[str, torch.Tensor]],
-        Dict[torch.Tensor],
+        Dict[str, torch.Tensor],
     ]:
         r"""Call Equi2Cube