Using MSE Loss Function in Classification

fastai fastai dev
1

I have been trying to create a function in which I train any pytorch model by its name using FastAI.

dls = DataLoaders(train_loader, valid_loader) #train_loader and valid_loader are pure pytorch loader
model = eval("torchvision.models.{}".format(model_name))

learner = vision_learner(
            dls,
            model,
            loss_func=loss_func,
            lr=lr,
            opt_func=optimizer,
            n_out=n_classes,
            pretrained=True,
            cut=cut,
            normalize=False,
        )

start = time.time()
learner.fit(n_epochs)
end = time.time()

When I set the n_classes variables (equal to n_out) as 1, the model works perfectly fine. However, changing it to another value gives the error:
The size of tensor a (4) must match the size of tensor b (2) at non-singleton dimension 0.

This occurs with MSELossFlat() and BCEWithLogitsLossFlat() but not with nn.CrossEntropyLoss()