I have a similar problem.
To get the gradient w.r.t input, I add the .requires_grad
but it doesn’t work.
for num,(sample_img, sample_label) in enumerate(mnist_test):
if num == 1:
break
sample_img = sample_img.to(device)
sample_img.requires_grad = True
prediction = model(sample_img.unsqueeze(dim=0))
cost = criterion(prediction, torch.tensor([sample_label]).to(device))
optimizer.zero_grad()
cost.backward()
print(sample_label)
print(sample_img.shape)

plt.imshow(sample_img.detach().cpu().squeeze(), cmap='gray')
plt.show()