I have a model with multiple outputs. How to use multiple losses?
import torch.nn as nn class NeuralNetwork(nn.Module): def __init__(self): super(NeuralNetwork, self).__init__() self.linear1 = nn.Linear(in_features = 3, out_features = 1) self.linear2 = nn.Linear(in_features = 3,out_features = 2) def forward(self, x): output1 = self.linear1(x) output2 = self.linear2(x) return output1, output2
output1 - CrossEntropyLoss, output2 - MSE
For one output it is easy:
criterion = nn.CrossEntropyLoss() trainer = create_supervised_trainer(model, optimizer, criterion, device=device)