add type annotations to torch.nn.parallel._functions

mypy.ini

@@ -86,9 +86,6 @@ ignore_errors = True
 [mypy-torch.nn.modules.pooling]
 ignore_errors = True
 
-[mypy-torch.nn.parallel._functions]
-ignore_errors = True
-
 [mypy-torch.nn.qat.modules.activations]
 ignore_errors = True
 

torch/nn/parallel/_functions.py

@@ -4,6 +4,7 @@
 from . import comm
 from torch.autograd import Function
 from torch._utils import _get_device_index
+from typing import List, Optional
 
 
 class Broadcast(Function):
@@ -39,9 +40,9 @@ class ReduceAddCoalesced(Function):
     def forward(ctx, destination, num_inputs, *grads):
         ctx.target_gpus = [grads[i].get_device() for i in range(0, len(grads), num_inputs)]
 
-        grads = [grads[i:i + num_inputs]
-                 for i in range(0, len(grads), num_inputs)]
-        return comm.reduce_add_coalesced(grads, destination)
+        grads_ = [grads[i:i + num_inputs]
+                  for i in range(0, len(grads), num_inputs)]
+        return comm.reduce_add_coalesced(grads_, destination)
 
     @staticmethod
     def backward(ctx, *grad_outputs):
@@ -105,10 +106,10 @@ def backward(ctx, *grad_output):
 
 
 # background streams used for copying
-_streams = None
+_streams: Optional[List[Optional[torch.cuda.Stream]]] = None
 
 
-def _get_stream(device):
+def _get_stream(device: int):
     """Gets a background stream for copying between CPU and GPU"""
     global _streams
     if device == -1: