diff --git a/combo/models/model.py b/combo/models/model.py index cb386fb739a02eb6440145f07150a6cde4cd18d8..883934eb5e041c8eb94a107ea21bbd89da5b6bbf 100644 --- a/combo/models/model.py +++ b/combo/models/model.py @@ -53,7 +53,7 @@ class NERModel(heads.Head): ) if tags is not None: - self._loss = output["loss"] + self._loss = output["loss"].cpu().item() self._accuracy_metric(output["probability"], tags, word_mask) self._f1_metric(output["probability"], tags, word_mask)