Gradient descent on the input to a model trained to recognize MNIST digits

GD_on_input_MNIST.py

x = torch.ones(bs, 784) * 0.5 + 0.02 * torch.randn(bs, 784)
x.requires_grad_(True).retain_grad()
for i in range(0, n//bs):
    preds = model(x)
    loss = loss_func(preds, y)
    loss.backward()
    with torch.no_grad():
        x -= x.grad * lr
        if x.grad.grad_fn is not None:
            x.grad.detach_()
        else:
            x.grad.requires_grad_(False)
        x.grad.zero_()
        model.zero_grad()
report(loss, preds, yb)

i = 20
plt.matshow(x[i].detach().numpy().reshape(-1, 28))
plt.matshow(xb[i].detach().numpy().reshape(-1, 28)); yb[i]