Fix DeferredBatchNorm mypy error

kakaobrain · Jun 7, 2019 · 6a0a4d0 · 6a0a4d0
1 parent bf51a06
commit 6a0a4d0
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 6 deletions.
diff --git a/stubs/torch/nn/modules/batchnorm.pyi b/stubs/torch/nn/modules/batchnorm.pyi
@@ -18,11 +18,6 @@ class _BatchNorm(Module):
     running_var: Tensor
     num_batches_tracked: Tensor
 
-    # for DeferredBatchNorm
-    sum: Tensor
-    sum_squares: Tensor
-    counter: Tensor
-
     def __init__(self,
                  num_features: int,
                  eps: float = 1e-5,

diff --git a/torchgpipe/batchnorm.py b/torchgpipe/batchnorm.py
@@ -52,7 +52,7 @@ def _track(self, input: Tensor) -> bool:
             self.sum += input.sum(dim)
             self.sum_squares += (input**2).sum(dim)
 
-        size = input.size().numel() / input.size(1)
+        size = input.size().numel() // input.size(1)
         self.counter += size
         self.tracked += 1