ModMSELoss is incorrect

[Cogito2012]

Thanks for your great work! But there seems to be an incorrect implementation for the ModMSELoss.

1. The division operation (output / output_max) could cause NaN values if all entries of output are zeros.
2. The square operation should take the whole fraction term as input.

Therefore, I changed the forward function of ModMSELoss as follows:

def forward(self, output , label , prior):
        prior_size = prior.shape
        output_max = torch.max(torch.max(output,2)[0],2)[0].unsqueeze(2).unsqueeze(2).expand(output.shape[0],output.shape[1],self.shape_r_gt,self.shape_c_gt)
        reg = ( 1.0/(prior_size[0]*prior_size[1]) ) * ( 1 - prior)**2  # (1, 1, 6, 8)
        loss = torch.mean( ((output / (output_max + 1e-6) - label) / (1 - label + 0.1))**2)  +  torch.sum(reg)
        return loss