Solved by ptrblck in post #5 
              
              
                This should create the same gradients. 
I’ve used this simple dummy snippet to verify it: 
# Setup
torch.manual_seed(2809)

model = nn.Sequential(
    nn.Conv2d(3, 6, 3, 1, 1),
    nn.ReLU(),
    nn.Conv2d(6, 1, 3, 1, 1),
    nn.ReLU(),
    nn.Flatten(),
    nn.Linear(1*24*24, 10)
)

x = torch.randn…