code: def __init__(
      self,
      sz: int,
      eps=1e-5,
      affine=True,
      track_running_stats=True,
      momentum=0.1,
  ):
      self.eps, self.track_running_stats, self.momentum = (
          eps,
          track_running_stats,
          momentum,
      )
  
      if affine:
          self.weight, self.bias = Tensor.ones(sz), Tensor.zeros(sz)
      else:
          self.weight, self.bias = None, None
  
      self.running_mean, self.running_var = Tensor.zeros(
          sz, requires_grad=False
      ), Tensor.ones(sz, requires_grad=False)
      self.num_batches_tracked = Tensor.zeros(1, requires_grad=False)
  
imports: import math
  from typing import Optional, Union, Tuple, cast
  from tinygrad.tensor import Tensor
  from tinygrad.helpers import prod
  from tinygrad.nn import optim, state
label_desc: Use image node_id 1 to create a new BatchNorm2d object  with example: obj = BatchNorm2d(sz, eps, affine, track_running_stats, momentum)