How to average loss whilst retaining grad_fn information?

I am trying to average the output of the nn.MSELoss() over the following batch before firing batch_loss.backward().

[tensor(0.0031, device='cuda:0', grad_fn=<MseLossBackward>),
 tensor(0.0010, device='cuda:0', grad_fn=<MseLossBackward>),
 tensor(0.0015, device='cuda:0', grad_fn=<MseLossBackward>),
 tensor(0.0027, device='cuda:0', grad_fn=<MseLossBackward>),
 tensor(0.0042, device='cuda:0', grad_fn=<MseLossBackward>),
 tensor(0.0045, device='cuda:0', grad_fn=<MseLossBackward>),
 tensor(0.0044, device='cuda:0', grad_fn=<MseLossBackward>),
 tensor(0.0005, device='cuda:0', grad_fn=<MseLossBackward>),
 tensor(0.0034, device='cuda:0', grad_fn=<MseLossBackward>),
 tensor(0.0023, device='cuda:0', grad_fn=<MseLossBackward>)]

How can I do this whilst retaining the grad_fn information? As if I use torch.mean(torch.tensor()), I lose this information.

Thanks!

sum(output).backward(). It will work if output is a list.

1 Like