def METHOD_NAME(
    self,
    x: torch.Tensor,
    label: torch.Tensor,
    label_mask: torch.BoolTensor = None,
    position_encoding: torch.Tensor = None,
    attention_mask: torch.LongTensor = None,
    unique_name: List[str] = None,
):
    loss, hidden_states, prediction = self.predict(
        x, label, label_mask, position_encoding, attention_mask
    ).slice(3)
    logs = Logs()
    logs.add_hidden_state("hidden_states", hidden_states)
    logs.add_hidden_state("prediction", prediction)
    return Output(
        loss=loss,
        prediction=prediction,
        label=label,
        unique_name=unique_name,
        logs=logs,
    )