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)