Config mypy correctly for torchgpipe_balancing

setup.cfg

@@ -55,7 +55,7 @@ mypy_path = ./stubs/
 follow_imports = normal
 
 # This project must be strictly typed.
-[mypy-torchgpipe.*,mypy-torchgpipe_balancing.*]
+[mypy-torchgpipe.*,torchgpipe_balancing.*]
 check_untyped_defs = true
 disallow_untyped_defs = true
 disallow_untyped_calls = true

torchgpipe_balancing/blockpartition.py

@@ -3,7 +3,7 @@
 Paper: https://arxiv.org/pdf/1308.2452.pdf
 
 """
-from typing import List
+from typing import Iterator, List, Tuple
 
 __all__ = ['solve']
 
@@ -40,7 +40,7 @@ def block_size(i: int) -> float:
         stop = splits[i]
         return sum(normal_sequence[start:stop])
 
-    def leaderboard():
+    def leaderboard() -> Iterator[Tuple[float, int]]:
         return ((block_size(i), i) for i in range(partitions))
 
     while True:

torchgpipe_balancing/utils.py

@@ -60,7 +60,7 @@ def training_sandbox(module: nn.Sequential) -> Generator[None, None, None]:
     module.train(training)
 
 
-def synchronize_device(device: torch.device):
+def synchronize_device(device: torch.device) -> None:
     if device.type == 'cpu':
         return
     torch.cuda.synchronize(device)