Fix mypy errors

kakaobrain · Aug 16, 2019 · 8e7b8bd · 8e7b8bd
1 parent 310c2a6
commit 8e7b8bd
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/stubs/torch/cuda/__init__.pyi b/stubs/torch/cuda/__init__.pyi
@@ -26,7 +26,7 @@ _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: ...
+def synchronize(device: _device_t = ...) -> None: ...
 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: ...

diff --git a/stubs/torch/nn/modules/module.pyi b/stubs/torch/nn/modules/module.pyi
@@ -95,6 +95,10 @@ class Module(Generic[T_co]):
     def named_modules(self, memo: Optional[Set['Module']] = ..., prefix: str = ...) -> Iterator[
         Tuple[str, 'Module']]: ...
 
+#MODIFIED BY TORCHGPIPE
+    training: bool
+#END
+
     def train(self: T, mode: bool = ...) -> T: ...
 
     def eval(self: T) -> T: ...