def forward(self, tensor, labels=None):
    logits = self.model(tensor)
    if labels is not None:
        loss = torch.nn.cross_entropy(logits, labels)
        return {"loss": loss, "logits": logits}
    return {"logits": logits}