def METHOD_NAME(self, inputs):
    image, label = inputs
    pred = self.model(image.to(self.device))
    if self.training:
        label = label.to(self.device)
        return F.cross_entropy(pred, label)
    else:
        return pred