And fix pyi

malfet · malfet · commit df7888e2dafc · 2023-10-25T18:13:13.000-07:00
diff --git a/tools/pyi/gen_pyi.py b/tools/pyi/gen_pyi.py
@@ -983,12 +983,12 @@ def replace_special_case(hint: str) -> str:
     unsorted_tensor_method_hints.update(
         {
             "size": [
-                "def size(self) -> Size: ...",
+                "def size(self, dim: None = None) -> Size: ...",
                 "def size(self, dim: _int) -> _int: ...",
             ],
             "stride": [
-                "def stride(self) -> Tuple[_int, ...]: ...",
-                "def stride(self, _int) -> _int: ...",
+                "def stride(self, dim: None = None) -> Tuple[_int, ...]: ...",
+                "def stride(self, dim: _int) -> _int: ...",
             ],
             "new_ones": [
                 f"def new_ones(self, size: _size, {FACTORY_PARAMS}) -> Tensor: ..."
