Hi,
Your error comes from the fact that one of your lambda returns a tensor which is not of the same size as what it got as input. You can replace the lambda by a real function and add some extra prints. In this case, I think the batch size is handled in a way you don’t expect: the gradient of the hook for a tensor of size HxW will be BxHxW and is then summed over the batch dimension if I remember correctly.
For your use case, if you just want to set the gradients, you don’t actually need to run a backward pass:
optimizer.zero_grad()
with torch.no_grad():
  for (key, w) in model.named_parameters():
    w.grad.copy_(new_gradients[key])

optimizer.step()