update

loss_metric.py

@@ -15,8 +15,10 @@
 
 import evaluate
 import datasets
+import numpy as np
+import torch
 
-from torch import nn, Tensor
+from torch import nn, Tensor, tensor
 
 _CITATION = """\
 @InProceedings{huggingface:module,
@@ -75,11 +77,12 @@ class LossMetric(evaluate.Metric):
 
     def _compute(self, predictions, references):
         """Returns the scores"""
-        print(predictions)
+        pred = tensor(np.array(predictions), dtype=torch.float16)
+        label = tensor(np.array(references), dtype=torch.float16)
         loss_func = nn.CrossEntropyLoss()
-        loss = loss_func(predictions, references)
+        loss = loss_func(pred, label)
 
-        mean_loss = loss.item() / references.shape[0]
+        mean_loss = loss.item() / label.shape[0]
         return {
             "loss": mean_loss,
         }