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!