Classification
Single Label
loss_fn = nn.CrossEntropyLoss()
# Forward pass
logits = model(inputs)
# Backward pass
optimizer.zero_grad()
loss = loss_fn(logits, targets)
loss.backward()
optimizer.step()
Multi Label
loss_fn = nn.BCEWithLogitsLoss() # Binary Cross Entropy
# Forward pass
logits = model(inputs)
# Backward pass
optimizer.zero_grad()
loss = loss_fn(logits, targets)
loss.backward()
optimizer.step()