Memory optimization of perplexity metric (#2346)

* reduce memory footprint when computing perplexity

(cherry picked from commit c0d2d3a)

CHANGELOG.md

@@ -27,6 +27,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Fixed high memory consumption in `Perplexity` metric ([#2346](https://github.com/Lightning-AI/torchmetrics/pull/2346))
+
+
 - Fixed cached network in `FeatureShare` not being moved to the correct device ([#2348](https://github.com/Lightning-AI/torchmetrics/pull/2348))
 
 

src/torchmetrics/functional/text/perplexity.py

@@ -91,7 +91,7 @@ def _perplexity_update(preds: Tensor, target: Tensor, ignore_index: Optional[int
     else:
         mask = torch.ones_like(target, dtype=torch.bool)
 
-    probs = probs[:, target].diagonal()[mask]
+    probs = probs[torch.arange(target.numel()), target][mask]
     total_log_probs = -probs.log().sum()
     count = mask.sum()