code: def test_linear_mnist(self):
      class LinTiny:
          def __init__(self, has_batchnorm=False):
              self.l1 = Linear(784, 128)
              self.l2 = Linear(128, 10)
              self.bn1 = (
                  BatchNorm2d(128) if has_batchnorm else lambda x: x
              )
  
          def __call__(self, x):
              return self.l2(self.l1(x)).relu().log_softmax(-1)
  
      class LinTorch(nn.Module):
          def __init__(self, has_batchnorm=False):
              super().__init__()
              self.l1 = nn.Linear(784, 128)
              self.l2 = nn.Linear(128, 10)
  
          def forward(self, x):
              return self.l2(self.l1(x)).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 3 for calling the TestEnd2End obj's underlying member method code with example usage: obj.test_linear_mnist() and returns: self, self