Solved by ptrblck in post #2 
              
              
                Yes, the workflow should work as shown in this dummy example: 
baseline = nn.Linear(1, 1)
model = nn.Linear(1, 1)

for param in model.parameters():
    param.requires_grad_(False)

x = torch.randn(1, 1)
out = baseline(x)
print(out.grad_fn) # prints valid grad_fn
out = model(out)
out.mean().backward(…