eval_model peut recevoir des modèles huggingface

evaluation.py

@@ -13,7 +13,11 @@ def eval_model(model, dataloader, device):
         labels = labels.to(device)
 
         logits = model(x, attention_mask)
+        try:
             predicted = torch.argmax(logits, dim=-1).tolist()
+        except TypeError:
+            # Si le modèle provient de huggingface
+            predicted = torch.argmax(logits["logits"], dim=-1).tolist()
 
         for pred, gold in zip(predicted, labels.tolist()):
             if gold not in confusion: