def METHOD_NAME(
    self, batch: TorchData, epoch_idx: int, batch_idx: int
) -> Dict[str, torch.Tensor]:
    batch = cast(Tuple[torch.Tensor, torch.Tensor], batch)
    data, labels = batch
    output = self.model(data)
    loss = torch.nn.functional.nll_loss(output, labels)
    # Update our custom metric.
    self.f1_score.update(y_true=labels, y_pred=output.argmax(dim=1))
    self.context.backward(loss)
    self.context.step_optimizer(self.optimizer)
    return {"loss": loss}