Solved by ptrblck in post #8 
              
              
                I’m not sure, why the shapes differ, but apparently the wrong gradients are stored. 
Here is a small dummy example using vgg16: 
grads = []
def save_grad(grad):
    grads.append(grad)

# Create model
model = models.vgg16()
model.eval()

# First approach
x = torch.randn(1, 3, 224, 224)
output = model…