def METHOD_NAME(self, batch: TorchData, epoch_idx: int, batch_idx: int):
    """
    Trains the provided batch.
    Returns: Dictionary of the calculated Metrics
    """
    inputs = {"input_ids": batch[0], "attention_mask": batch[1], "labels": batch[3]}
    if self.context.get_hparam("model_type") != "distilbert":
        inputs["token_type_ids"] = (
            batch[2] if self.context.get_hparam("model_type") in ["bert", "xlnet"] else None
        )
    outputs = self.model(**inputs)
    results = self.get_metrics(outputs, inputs)
    self.context.backward(results["loss"])
    self.context.step_optimizer(self.optimizer)
    return results