Pytorch - vanilla Gradient Visualization
I trained a neural network on MNIST using PyTorch:
class MnistCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d( 1, 16, 3, stride = 1, padding = 2)
        self.pool1 = nn.MaxPool2d(kernel_size = 2)
        self.conv2 = nn.Conv2d(16, 32, 3, stride = 1, padding = 2)
        self.pool2 = nn.MaxPool2d(kernel_size = 2)
        self.dropout = nn.Dropout(0.5)
        self.lin = nn.Linear(32 * 8 * 8, 10)
    
    def forward(self, x):
        # conv block
        x = F.relu(self.conv1(x))
        x = self.pool1(x)
        # conv block
        x = F.relu(self.conv2(x))
        x = self.pool2(x)
        # dense block
        x = x.view(x.size(0), -1)
        x = self.dropout(x)
        return self.lin(x)

I would like to implement vanilla Gradient Visualization (see reference below) on my model.
Simonyan, K., Vedaldi, A., Zisserman, A.
Deep inside convolutional networks: Visualising image classification models and saliency maps.
arXiv preprint arXiv:1312.6034 (2013)

Question: How can I implement this method in PyTorch?
If I understand correctly, vanilla gradient visualization consists in computing the partial derivatives of the loss of my model w.r.t all the pixels in my input image. So to make it short, I need to tweek my self.conv1 layer so that it computes the gradient over its input pixels instead of the gradient over its weights.
Please correct me if I'm wrong.

You do not need to change anything about your conv layer. Each layer computes gradients both w.r.t. parameters (for updates) and w.r.t. inputs (for &quot;downstream&quot; gradients by the chain rule). Therefore, all you need is to set your input image's x gradient property to true:
x, y = ... # get one image from MNIST
x.requires_grad_(True)  # indicate to pytorch that you would like to look at these gradients
pred = model(x)
loss = criterion(pred, y)
loss.backward()  # propagate gradients
x.grad  # &lt;- here you should have the gradients of the loss w.r.t pixels

