update

loss_metric.py

@@ -16,7 +16,7 @@
 import evaluate
 import datasets
 
-from torch import nn
+from torch import nn, Tensor
 
 _CITATION = """\
 @InProceedings{huggingface:module,
@@ -55,7 +55,6 @@ class loss_metric(evaluate.Metric):
     """Calculation of the cross-entropy loss function using the huggingface evaluate module."""
 
     def _info(self):
-        # TODO: Specifies the evaluate.EvaluationModuleInfo object
         return evaluate.MetricInfo(
             # This is the description that will appear on the modules page.
             module_type="metric",
@@ -74,14 +73,12 @@ class loss_metric(evaluate.Metric):
             reference_urls=["http://path.to.reference.url/new_module"]
         )
 
-    def _download_and_prepare(self, dl_manager):
-        """Optional: download external resources useful to compute the scores"""
-        pass
-
-    def _compute(self, predictions, references):
+    def _compute(self, predictions: Tensor, references: Tensor):
         """Returns the scores"""
         loss_func = nn.CrossEntropyLoss()
         loss = loss_func(predictions, references)
+
+        mean_loss = loss.item() / references.shape[0]
         return {
-            "loss": loss,
+            "loss": mean_loss,
         }