code: def test_bn_mnist(self):
      class LinTiny:
          def __init__(self):
              self.l1 = Linear(784, 128)
              self.l2 = Linear(128, 10)
              self.bn1 = BatchNorm2d(128)
  
          def __call__(self, x):
              return self.l2(
                  self.bn1(self.l1(x).reshape(x.shape[0], -1, 1, 1))
                  .reshape(x.shape[0], -1)
                  .relu()
              ).log_softmax(-1)
  
      class LinTorch(nn.Module):
          def __init__(self):
              super().__init__()
              self.l1 = nn.Linear(784, 128)
              self.l2 = nn.Linear(128, 10)
              self.bn1 = nn.BatchNorm2d(128)
  
          def forward(self, x):
              return self.l2(
                  self.bn1(self.l1(x).reshape(x.shape[0], -1, 1, 1))
                  .reshape(x.shape[0], -1)
                  .relu()
              ).log_softmax(-1)
  
      compare_tiny_torch(LinTiny(), LinTorch(), self.X, self.Y)
  
file_path: tinygrad/test/models/test_end2end.py
imports: import torch
  from torch import nn
  import unittest
  import numpy
  from tinygrad.nn.state import get_parameters, get_state_dict
  from tinygrad.nn import optim, Linear, Conv2d, BatchNorm2d
  from tinygrad.tensor import Tensor
  from extra.datasets import fetch_mnist
  from tinygrad.helpers import CI
label_desc: Use image node_id 4 for calling the TestEnd2End obj's underlying member method code with example usage: obj.test_bn_mnist() and returns: self, self