@albanD Just wanted to let you know that this had to do with the shape of the tensors. We were able to fix the problem by using the MSELossFlat loss function provided in the fastai framework. This essentially wraps the torch.nn.MSELoss function in a flattener.
Below is a link to the fastai version of the loss function.

  
      github.com
  
  
    fastai/fastai/blob/8013797e05f0ae0d771d60ecf7cf524da591503c/fastai/layers.py#L257 32

    return FlattenedLoss(nn.CrossEntropyLoss, *args, axis=axis, **kwargs)


def BCEWithLogitsFlat(*args, axis:int=-1, floatify:bool=True, **kwargs):
    "Same as `nn.BCEWithLogitsLoss`, but flattens input and target."
    return FlattenedLoss(nn.BCEWithLogitsLoss, *args, axis=axis, floatify=floatify, is_2d=False, **kwargs)


def BCEFlat(*args, axis:int=-1, floatify:bool=True, **kwargs):
    "Same as `nn.BCELoss`, but flattens input and target."
    return FlattenedLoss(nn.BCELoss, *args, axis=axis, floatify=floatify, is_2d=False, **kwargs)


def MSELossFlat(*args, axis:int=-1, floatify:bool=True, **kwargs):
    "Same as `nn.MSELoss`, but flattens input and target."
    return FlattenedLoss(nn.MSELoss, *args, axis=axis, floatify=floatify, is_2d=False, **kwargs)


class NoopLoss(Module):
    "Just returns the mean of the `output`."
    def forward(self, output, *args): return output.mean()


class WassersteinLoss(Module):
    "For WGAN."
    def forward(self, real, fake): return real.mean() - fake.mean()



  
  
    
    
  
  


This an interesting bug because we have not encountered it before in our classification problems. We freely interchange pytorch and fastai loss functions, however, in this case it was a problem. We were working the rossmann sales problem.