def METHOD_NAME(self, batch: Any, batch_idx: int, metrics: nn.ModuleDict) -> Any:
    out = self.model(batch["input_values"], labels=batch["labels"])
    out["logs"] = {"loss": out.loss}
    return out