Method test_mnist_backward

test/test_tiny.py:135–146

(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_backward(self):
    # NOTE: we don't have the whole model here for speed
    layers = [
      nn.Conv2d(1, 8, 5), Tensor.relu,
      nn.Conv2d(8, 8, 5), Tensor.relu]

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

    # realize gradients
    Tensor.empty(4, 1, 14, 14).sequential(layers).sum().backward()
    Tensor.realize(*[x.grad for x in nn.state.get_parameters(layers) if x.grad is not None])

  # *** image ***