Write more type hints for PyTorch

stubs/torch/cuda/__init__.pyi

@@ -26,7 +26,11 @@ _device_t = Union[_device, int]
 def check_error(res: int) -> None: ...
 def device_count() -> int: ...
 def empty_cache() -> None: ...
-def synchronize(device: _device_t) -> None: ...
+
+#MODIFIED BY TORCHGPIPE
+def synchronize(device: Optional[_device_t] = ...) -> None: ...
+#END
+
 def set_device(device: _device_t) -> None: ...
 def get_device_capability(device: Optional[_device_t]=...) -> Tuple[int, int]: ...
 def get_device_name(device: Optional[_device_t]=...) -> str: ...
@@ -40,3 +44,11 @@ def max_memory_cached(device: Optional[_device_t]=...) -> int: ...
 def reset_max_memory_cached(device: Optional[_device_t]=...) -> None: ...
 def cudart() -> ctypes.CDLL: ...
 def find_cuda_windows_lib() -> Optional[ctypes.CDLL]: ...
+
+#MODIFIED BY TORCHGPIPE
+from typing import Any
+class device:
+    def __init__(self, device: _device_t = ...) -> None: ...
+    def __enter__(self) -> None: ...
+    def __exit__(self, *args: Any) -> None: ...
+#END

stubs/torch/nn/__init__.pyi

@@ -43,6 +43,15 @@ class Module:
     def named_children(self) -> Iterable[Tuple[str, Module]]: ...
     def add_module(self, name: str, module: Module) -> None: ...
 
+    def buffers(self) -> Iterator[Tensor]: ...
+    def parameters(self) -> Iterator[Parameter]: ...
+
+    def state_dict(self, destination: Optional[str] = ..., prefix: str = ..., keep_vars: bool = ...) -> Dict[str, Tensor]: ...
+    def load_state_dict(self, state_dict: Dict[str, Tensor], strict: bool = ...) -> Tuple[List[str], List[str]]: ...
+
+    def train(self: TModule, mode: bool = ...) -> TModule: ...
+    def eval(self: TModule) -> TModule: ...
+
 
 class Sequential(Module):
     @overload