def _forward(self, model_inputs):
    # model_inputs == {"model_input": model_input}
    outputs = self.model(**model_inputs)
    # Maybe {"logits": Tensor()}
    return outputs

def postprocess(self, model_outputs):
    best_class = model_outputs["logits"].softmax(-1)
    return best_class