Module.to and apply returns itself

kakaobrain · Jun 13, 2019 · bd1b0a8 · bd1b0a8
1 parent 353a302
commit bd1b0a8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/stubs/torch/nn/__init__.pyi b/stubs/torch/nn/__init__.pyi
@@ -30,8 +30,8 @@ class Module:
     def __init__(self) -> None: ...
     def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
 
-    def to(self, device: device) -> Module: ...
-    def apply(self, fn: Callable[[Module], None]) -> Module: ...
+    def to(self: TModule, device: device) -> TModule: ...
+    def apply(self: TModule, fn: Callable[[Module], None]) -> TModule: ...
 
     def register_buffer(self, name: str, tensor: Tensor) -> None: ...
     def register_parameter(self, name: str, param: Union[Parameter, None]) -> None: ...

diff --git a/torchgpipe/gpipe.py b/torchgpipe/gpipe.py
@@ -177,7 +177,7 @@ def cuda(self, device: Any = None) -> nn.Module:
     def cpu(self) -> nn.Module:
         raise MOVING_DENIED
 
-    def to(self, *args: Any, **kwargs: Any) -> nn.Module:
+    def to(self, *args: Any, **kwargs: Any) -> 'GPipe':
         # Deny these usages:
         #
         # - to(device[, dtype, non_blocking])