Add torch.squeeze parameter description to declare allowed type

zeshengzong · zeshengzong · commit bfebfd908d8f · 2024-10-08T16:10:44.000+08:00
diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py
@@ -10378,7 +10378,7 @@ def merge_dicts(*dicts):
 add_docstr(
     torch.squeeze,
     r"""
-squeeze(input, dim=None) -> Tensor
+squeeze(input: Tensor, dim: Optional[Union[int, List[int]]]) -> Tensor
 
 Returns a tensor with all specified dimensions of :attr:`input` of size `1` removed.
 
