code: def test_conv_mnist(self):
      class LinTiny:
          def __init__(self, has_batchnorm=False):
              self.c1 = Conv2d(1, 8, 3, stride=2)
              self.c2 = Conv2d(8, 16, 3, stride=2)
              self.l1 = Linear(16 * 6 * 6, 10)
              if has_batchnorm:
                  self.bn1, self.bn2 = BatchNorm2d(8), BatchNorm2d(16)
              else:
                  self.bn1, self.bn2 = lambda x: x, lambda x: x
  
          def __call__(self, x):
              return self.l1(
                  self.bn2(self.c2(self.bn1(self.c1(x)).relu()))
                  .relu()
                  .reshape(x.shape[0], -1)
              ).log_softmax(-1)
  
      class LinTorch(nn.Module):
          def __init__(self, has_batchnorm=False):
              super().__init__()
              self.c1 = nn.Conv2d(1, 8, 3, stride=2)
              self.c2 = nn.Conv2d(8, 16, 3, stride=2)
              self.l1 = nn.Linear(16 * 6 * 6, 10)
              if has_batchnorm:
                  self.bn1, self.bn2 = nn.BatchNorm2d(
                      8
                  ), nn.BatchNorm2d(16)
              else:
                  self.bn1, self.bn2 = lambda x: x, lambda x: x
  
          def forward(self, x):
              return self.l1(
                  self.bn2(self.c2(self.bn1(self.c1(x)).relu()))
                  .relu()
                  .reshape(x.shape[0], -1)
              ).log_softmax(-1)
  
      for has_batchnorm in [False, True]:
          with self.subTest(has_batchnorm=has_batchnorm):
              compare_tiny_torch(
                  LinTiny(has_batchnorm),
                  LinTorch(has_batchnorm),
                  self.X.reshape((-1, 1, 28, 28)),
                  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 7 for calling the TestEnd2End obj's underlying member method code with example usage: obj.test_conv_mnist() and returns: self, self