bn: don't track_running_stats if you aren't training
MODIFY
tinygrad/nn.py
tinygrad/nn.py
@@ -12,7 +12,7 @@ class BatchNorm2D:
self.num_batches_tracked = Tensor.zeros(1, requires_grad=False)
def __call__(self, x):
- if self.track_running_stats or Tensor.training:
+ if Tensor.training:
x_detached = x.detach()
batch_mean = x_detached.mean(axis=(0,2,3))
y = (x_detached - batch_mean.reshape(shape=[1, -1, 1, 1]))
@@ -24,7 +24,6 @@ class BatchNorm2D:
if self.num_batches_tracked is None: self.num_batches_tracked = Tensor.zeros(1, requires_grad=False)
self.num_batches_tracked += 1
- if Tensor.training:
return self.normalize(x, batch_mean, batch_var)
return self.normalize(x, self.running_mean, self.running_var)

bn: don't track_running_stats if you aren't training