def METHOD_NAME(self, model, inputs):
    output = model(*inputs)
    output = output.mean()
    if isinstance(output, (list, tuple)):
        loss = output[0]
    else:
        loss = output
    return loss