[BE]: Improve typing in torch/modules/container.py

Skylion007 · Skylion007 · commit 16c461947528 · 2025-05-16T11:32:21.000-04:00
diff --git a/torch/nn/modules/container.py b/torch/nn/modules/container.py
@@ -121,7 +121,7 @@ def __init__(self, *args):
             for idx, module in enumerate(args):
                 self.add_module(str(idx), module)
 
-    def _get_item_by_idx(self, iterator, idx) -> T:  # type: ignore[misc, type-var]
+    def _get_item_by_idx(self, iterator: Iterable[T], idx: int) -> T:  # type: ignore[misc, type-var]
         """Get the idx-th item of the iterator."""
         size = len(self)
         idx = operator.index(idx)
@@ -298,7 +298,7 @@ def insert(self, index: int, module: Module) -> Self:
         self._modules[str(index)] = module
         return self
 
-    def extend(self, sequential) -> Self:
+    def extend(self, sequential: Iterable[Module]) -> Self:
         """
         Extends the current Sequential container with layers from another Sequential container.
 
