Update app.py

app.py

@@ -82,7 +82,7 @@ def train_model(dataset):
 
     # Initialize the Trainer with custom compute_loss to verify shapes
     class CustomTrainer(Trainer):
-        def compute_loss(self, model, inputs, return_outputs=False):
+        def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
             outputs = model(**inputs)
             logits = outputs.logits
             labels = inputs['labels']