nn.MSELoss inputs inquiry

I have a question with regards to nn.MSELoss(). I’m trying to implement a network to demosaic/debayer an image. I have seen that the inputs to nn.MSELoss should be nn.MSELoss(outputs, target). My outputs tensor is of shape [batch_size, 3, 224, 224], but my target is [batch_size].

I am loading the images as so.

trainset = datasets.ImageFolder(root=trainpath,transform=preprocess)
trainloader = DataLoader(dataset=trainset,batch_size=batch_size, shuffle=True)

And then in my training function

for data in trainloader:
    inputs, target = data[0], data[1]
    outputs = model(inputs)
    loss = nn.MSELoss(outputs, target)

Should it instead be loss = nn.MSELoss(outputs, inputs) as I know MSE subtracts the two values. But from everything I’ve seen, it’s always (outputs, target).

Thank you very much

As you’ve explained, nn.MSELoss expects the model output and target to have the same shape.
Since your target shape differs, you would have to explain a bit more how the loss calculation should be performed for your use case.
E.g. you could repeat the target values for each channel and pixel, but I doubt this would fit your use case.