Update app.py

app.py

@@ -54,7 +54,7 @@ def train_model(dataset):
     )
 
     class CustomTrainer(Trainer):
-        def compute_loss(self, model, inputs, return_outputs=False):
+        def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
             labels = inputs.pop("labels")
             outputs = model(**inputs)
             logits = outputs.logits