Hi, I’m new in the pytorch.
I have a question about the custom loss function.
The code is following.
I use numpy to clone the MSE_loss as MSE_SCORE.
Input is 1x200x200 images, and batch size is 128.
The output “mse” of the MSE_SCORE is a float value and covered as tensor.
However, when I tried to apply the backward of custom loss function, then got the error “RuntimeError: Function MSE_SCOREBackward returned an invalid gradient at index 0 - got [] but expected shape compatible with [128, 1, 200, 200]”
 class MSE_SCORE(Function):
     @staticmethod
     def forward(ctx, input, label):
         with torch.no_grad():
             numpy_input = input.detach().numpy()
             numpy_label = label.detach().numpy()
             offset_height = int(height / 4)
             offset_width = int(width / 4)
             target_height = int(height / 2)
             target_width = int(width / 2)
             crop_input = numpy_input[:, :, offset_height:offset_height + target_height, offset_width:offset_width + target_width]
             crop_label = numpy_label[:, :, offset_height:offset_height + target_height, offset_width:offset_width + target_width]
             diff = (crop_input - crop_label)**2
             mean = np.mean(diff, axis=(1,2,3))
             mse = np.sum(mean)
 
         return torch.tensor(mse)
 
     @staticmethod
     def backward(ctx, grad_output):
         return grad_output, None
 
 def criterion_mse(input, label):
     return MSE_SCORE.apply(input, label)
 
 for epoch in range(num_epochs):
     for step, (img, label) in enumerate(dataloader):
         img = Variable(img).to(device=device)
         label = Variable(label).to(device=device)
         output = model(img)
         mse_loss = criterion_mse(output.cpu(), label.cpu()).to(device=device)
         optimizer.zero_grad()
         loss_seq = [mse_loss]
         grad_seq = [torch.tensor(1.0).to(device=device) for _ in range(len(loss_seq))]
         torch.autograd.backward(loss_seq, grad_seq)