switch back to torchmetrics.text.Perplexity

NVIDIA · Dec 27, 2024 · 4c77f97 · 4c77f97
1 parent cb3de59
commit 4c77f97
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/sub-packages/bionemo-llm/src/bionemo/llm/lightning.py b/sub-packages/bionemo-llm/src/bionemo/llm/lightning.py
@@ -294,9 +294,11 @@ def __init__(
         self.log_train_ppl = log_train_ppl
         self.log_val_ppl = log_val_ppl
         if log_train_ppl:
-            self.train_ppl = MegatronPerplexityMetric(ignore_index=-100)
+            # self.train_ppl = MegatronPerplexityMetric(ignore_index=-100)
+            self.train_ppl = torchmetrics.text.Perplexity(ignore_index=-100)
         if log_val_ppl:
-            self.valid_ppl = MegatronPerplexityMetric(ignore_index=-100)
+            # self.valid_ppl = MegatronPerplexityMetric(ignore_index=-100)
+            self.valid_ppl = torchmetrics.text.Perplexity(ignore_index=-100)
 
     def configure_model(self) -> None:
         """Updates internal state: instantiates the model from the object's config, assigns to `model` attribute.