Method test_mnist

test/test_tiny.py:116–131

(self)

  # TODO: this is failing because of how swizzling rewrites the ShapeTracker of the final STORE
  @unittest.skipIf(CI and Device.DEFAULT == "DSP", "failing because of make things that can't be images not images")
  def test_mnist(self):
    layers = [
      nn.Conv2d(1, 32, 5), Tensor.relu,
      nn.Conv2d(32, 32, 5), Tensor.relu,
      nn.BatchNorm(32), Tensor.max_pool2d,
      nn.Conv2d(32, 64, 3), Tensor.relu,
      nn.Conv2d(64, 64, 3), Tensor.relu,
      nn.BatchNorm(64), Tensor.max_pool2d,
      lambda x: x.flatten(1), nn.Linear(576, 10)]

    # replace random weights with ones
    Tensor.realize(*[p.replace(Tensor.ones_like(p).contiguous()) for p in nn.state.get_parameters(layers)])

    # run model inference
    probs = Tensor.rand(1, 1, 28, 28).sequential(layers).tolist()
    self.assertEqual(len(probs[0]), 10)

  # TODO: this is failing because of how swizzling rewrites the ShapeTracker of the final STORE
  @unittest.skipIf(CI and Device.DEFAULT == "DSP", "failing because of make things that can't be images not images")