num_batches_tracked has shape ()

2025-10-21 09:10:35 +08:00
1 changed files with 1 additions and 1 deletions
@@ -36,7 +36,7 @@ class BatchNorm:
    self.weight: Tensor|None = Tensor.ones(sz) if affine else None
    self.bias: Tensor|None = Tensor.zeros(sz) if affine else None

-    self.num_batches_tracked = Tensor.zeros(1, dtype='long' if is_dtype_supported(dtypes.long) else 'int', requires_grad=False)
+    self.num_batches_tracked = Tensor.zeros(dtype='long' if is_dtype_supported(dtypes.long) else 'int', requires_grad=False)
    if track_running_stats: self.running_mean, self.running_var = Tensor.zeros(sz, requires_grad=False), Tensor.ones(sz, requires_grad=False)

  def calc_stats(self, x:Tensor) -> tuple[Tensor, Tensor]: