We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
High memory usage of Perplexity metric
## 🐛 Bug

I ran out of memory (GPU) when computing the perplexity metric and would like to propose a small optimization to decrease its memory utilization.

### To Reproduce

For instance, when running the following code PyTorch tries to allocate 1024 GB of GPU memory on my system.
```py
from torchmetrics.text import Perplexity
import torch

gen = torch.manual_seed(42)
preds = torch.rand(512, 1024, 12, generator=gen).cuda()
target = torch.randint(12, (512, 1024), generator=gen).cuda()

perp = Perplexity().cuda()
print(perp(preds, target))
```

### Memory Inefficiency

I think the inefficiency is in this line:

https://github.com/Lightning-AI/torchmetrics/blob/a68455afb9041d1d32c1d6546897fee416abdc41/src/torchmetrics/functional/text/perplexity.py#L94

`probs[:, target]` results in a large temporary tensor with `(512*1024)^2` elements. Afterwards only the diagonal values are used. 

### Potential Solution

In contrast

```
probs = probs[torch.arange(target.numel()), target][mask]
```

would only require memory of the size of target. 

Would you consider accepting a pull request with this optimization? Or was the previous implementation chosen for another reason?

### Environment

- TorchMetrics v1.2.1 (installed with pip) and Master branch.
- Python 3.10.12
- Pytorch 2.2.0
- CUDA 12.1

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/torchmetrics/functional/text/perplexity.py`
Content:
```
1 # Copyright The Lightning team.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 from typing import Optional, Tuple
16 
17 import torch
18 from torch import Tensor
19 
20 
21 def _check_shape_and_type_consistency(preds: Tensor, target: Tensor) -> None:
22     """Check shape and type consistency of input vectors.
23 
24     Args:
25         preds:
26             Logits or a unnormalized score assigned to each token in a sequence with shape [batch_size, seq_len,
27             vocab_size]. Scores will be normalized internally using softmax.
28         target:
29             Ground truth values with a shape [batch_size, seq_len].
30 
31     Raises:
32         ValueError:
33             If ``preds`` tensor has no 3 dimensions.
34         ValueError:
35             If ``target`` tensor has no 2 dimensions.
36         ValueError:
37             If the first two dimensions of ``preds`` and ``target`` do not equal.
38         TypeError:
39             If ``preds`` dtype is not one of ``(torch.float16, torch.float32, torch.float64)``
40         TypeError:
41             If ``target`` is not of a type LongTensor (torch.int64)
42 
43     """
44     if len(preds.shape) != 3:
45         raise ValueError(
46             "Input tensor `preds` is expected to have 3 dimensions, [batch_size, seq_len, vocab_size],"
47             f" but got {len(preds.shape)}."
48         )
49     if len(target.shape) != 2:
50         raise ValueError(
51             "Input tensor `target` is expected to have 2 dimensions, [batch_size, seq_len],"
52             f" but got {len(target.shape)}."
53         )
54     if preds.shape[:2] != target.shape:
55         raise ValueError(
56             "Input tensors `preds` and `target` are expected to have equaling first two dimensions,"
57             f" [batch_size, seq_len], but got {preds.shape[:2]} and {target.shape}."
58         )
59     if not preds.is_floating_point():
60         raise TypeError(f"Input tensor `preds` is expected to be of floating point type but got {preds.dtype}.")
61     if target.dtype != torch.int64:
62         raise TypeError(f"Input tensor `target` is expected to be of a type {torch.int64} but got {target.dtype}.")
63 
64 
65 def _perplexity_update(preds: Tensor, target: Tensor, ignore_index: Optional[int] = None) -> Tuple[Tensor, Tensor]:
66     """Compute intermediate statistics for Perplexity.
67 
68     Args:
69         preds:
70             Logits or a unnormalized score assigned to each token in a sequence with shape [batch_size, seq_len,
71             vocab_size]. Scores will be normalized internally using softmax.
72         target:
73             Ground truth values with a shape [batch_size, seq_len].
74         ignore_index:
75             Integer specifying a target class to ignore. If given, this class index does not contribute
76             to the returned score.
77 
78     Returns:
79         Log probabilities, summed over all samples
80         Number of samples
81 
82     """
83     _check_shape_and_type_consistency(preds, target)
84 
85     probs = torch.nn.functional.softmax(preds.reshape(-1, preds.shape[-1]), dim=1)
86     target = target.reshape(-1)
87 
88     if ignore_index is not None:
89         mask = target.ne(ignore_index)
90         target = target.where(target != ignore_index, torch.tensor(0, device=target.device))
91     else:
92         mask = torch.ones_like(target, dtype=torch.bool)
93 
94     probs = probs[:, target].diagonal()[mask]
95     total_log_probs = -probs.log().sum()
96     count = mask.sum()
97 
98     return total_log_probs, count
99 
100 
101 def _perplexity_compute(total: Tensor, count: Tensor) -> Tensor:
102     """Compute the Perplexity.
103 
104     Args:
105         total: Log probabilities, summed over all samples
106         count: Number of samples
107     Returns:
108         Perplexity
109 
110     """
111     return torch.exp(total / count)
112 
113 
114 def perplexity(preds: Tensor, target: Tensor, ignore_index: Optional[int] = None) -> Tensor:
115     """Perplexity measures how well a language model predicts a text sample.
116 
117     This metric is calculated as the average number of bits per word a model needs to represent the sample.
118 
119     Args:
120         preds:
121             Logits or a unnormalized score assigned to each token in a sequence with shape [batch_size, seq_len,
122             vocab_size], which is the output of a language model. Scores will be normalized internally using softmax.
123         target:
124             Ground truth values with a shape [batch_size, seq_len].
125         ignore_index:
126             Integer specifying a target class to ignore. If given, this class index does not contribute
127             to the returned score.
128 
129     Returns:
130         Perplexity value
131 
132     Examples:
133         >>> import torch
134         >>> gen = torch.manual_seed(42)
135         >>> preds = torch.rand(2, 8, 5, generator=gen)
136         >>> target = torch.randint(5, (2, 8), generator=gen)
137         >>> target[0, 6:] = -100
138         >>> perplexity(preds, target, ignore_index=-100)
139         tensor(5.8540)
140 
141     """
142     total, count = _perplexity_update(preds, target, ignore_index)
143     return _perplexity_compute(total, count)
144
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/src/torchmetrics/functional/text/perplexity.py b/src/torchmetrics/functional/text/perplexity.py
--- a/src/torchmetrics/functional/text/perplexity.py
+++ b/src/torchmetrics/functional/text/perplexity.py
@@ -91,7 +91,7 @@
     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()

{"golden_diff": "diff --git a/src/torchmetrics/functional/text/perplexity.py b/src/torchmetrics/functional/text/perplexity.py\n--- a/src/torchmetrics/functional/text/perplexity.py\n+++ b/src/torchmetrics/functional/text/perplexity.py\n@@ -91,7 +91,7 @@\n     else:\n         mask = torch.ones_like(target, dtype=torch.bool)\n \n-    probs = probs[:, target].diagonal()[mask]\n+    probs = probs[torch.arange(target.numel()), target][mask]\n     total_log_probs = -probs.log().sum()\n     count = mask.sum()\n", "issue": "High memory usage of Perplexity metric\n## \ud83d\udc1b Bug\r\n\r\nI ran out of memory (GPU) when computing the perplexity metric and would like to propose a small optimization to decrease its memory utilization.\r\n\r\n### To Reproduce\r\n\r\nFor instance, when running the following code PyTorch tries to allocate 1024 GB of GPU memory on my system.\r\n```py\r\nfrom torchmetrics.text import Perplexity\r\nimport torch\r\n\r\ngen = torch.manual_seed(42)\r\npreds = torch.rand(512, 1024, 12, generator=gen).cuda()\r\ntarget = torch.randint(12, (512, 1024), generator=gen).cuda()\r\n\r\nperp = Perplexity().cuda()\r\nprint(perp(preds, target))\r\n```\r\n\r\n### Memory Inefficiency\r\n\r\nI think the inefficiency is in this line:\r\n\r\nhttps://github.com/Lightning-AI/torchmetrics/blob/a68455afb9041d1d32c1d6546897fee416abdc41/src/torchmetrics/functional/text/perplexity.py#L94\r\n\r\n`probs[:, target]` results in a large temporary tensor with `(512*1024)^2` elements. Afterwards only the diagonal values are used. \r\n\r\n### Potential Solution\r\n\r\nIn contrast\r\n\r\n```\r\nprobs = probs[torch.arange(target.numel()), target][mask]\r\n```\r\n\r\nwould only require memory of the size of target. \r\n\r\nWould you consider accepting a pull request with this optimization? Or was the previous implementation chosen for another reason?\r\n\r\n### Environment\r\n\r\n- TorchMetrics v1.2.1 (installed with pip) and Master branch.\r\n- Python 3.10.12\r\n- Pytorch 2.2.0\r\n- CUDA 12.1\n", "before_files": [{"content": "# Copyright The Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Optional, Tuple\n\nimport torch\nfrom torch import Tensor\n\n\ndef _check_shape_and_type_consistency(preds: Tensor, target: Tensor) -> None:\n    \"\"\"Check shape and type consistency of input vectors.\n\n    Args:\n        preds:\n            Logits or a unnormalized score assigned to each token in a sequence with shape [batch_size, seq_len,\n            vocab_size]. Scores will be normalized internally using softmax.\n        target:\n            Ground truth values with a shape [batch_size, seq_len].\n\n    Raises:\n        ValueError:\n            If ``preds`` tensor has no 3 dimensions.\n        ValueError:\n            If ``target`` tensor has no 2 dimensions.\n        ValueError:\n            If the first two dimensions of ``preds`` and ``target`` do not equal.\n        TypeError:\n            If ``preds`` dtype is not one of ``(torch.float16, torch.float32, torch.float64)``\n        TypeError:\n            If ``target`` is not of a type LongTensor (torch.int64)\n\n    \"\"\"\n    if len(preds.shape) != 3:\n        raise ValueError(\n            \"Input tensor `preds` is expected to have 3 dimensions, [batch_size, seq_len, vocab_size],\"\n            f\" but got {len(preds.shape)}.\"\n        )\n    if len(target.shape) != 2:\n        raise ValueError(\n            \"Input tensor `target` is expected to have 2 dimensions, [batch_size, seq_len],\"\n            f\" but got {len(target.shape)}.\"\n        )\n    if preds.shape[:2] != target.shape:\n        raise ValueError(\n            \"Input tensors `preds` and `target` are expected to have equaling first two dimensions,\"\n            f\" [batch_size, seq_len], but got {preds.shape[:2]} and {target.shape}.\"\n        )\n    if not preds.is_floating_point():\n        raise TypeError(f\"Input tensor `preds` is expected to be of floating point type but got {preds.dtype}.\")\n    if target.dtype != torch.int64:\n        raise TypeError(f\"Input tensor `target` is expected to be of a type {torch.int64} but got {target.dtype}.\")\n\n\ndef _perplexity_update(preds: Tensor, target: Tensor, ignore_index: Optional[int] = None) -> Tuple[Tensor, Tensor]:\n    \"\"\"Compute intermediate statistics for Perplexity.\n\n    Args:\n        preds:\n            Logits or a unnormalized score assigned to each token in a sequence with shape [batch_size, seq_len,\n            vocab_size]. Scores will be normalized internally using softmax.\n        target:\n            Ground truth values with a shape [batch_size, seq_len].\n        ignore_index:\n            Integer specifying a target class to ignore. If given, this class index does not contribute\n            to the returned score.\n\n    Returns:\n        Log probabilities, summed over all samples\n        Number of samples\n\n    \"\"\"\n    _check_shape_and_type_consistency(preds, target)\n\n    probs = torch.nn.functional.softmax(preds.reshape(-1, preds.shape[-1]), dim=1)\n    target = target.reshape(-1)\n\n    if ignore_index is not None:\n        mask = target.ne(ignore_index)\n        target = target.where(target != ignore_index, torch.tensor(0, device=target.device))\n    else:\n        mask = torch.ones_like(target, dtype=torch.bool)\n\n    probs = probs[:, target].diagonal()[mask]\n    total_log_probs = -probs.log().sum()\n    count = mask.sum()\n\n    return total_log_probs, count\n\n\ndef _perplexity_compute(total: Tensor, count: Tensor) -> Tensor:\n    \"\"\"Compute the Perplexity.\n\n    Args:\n        total: Log probabilities, summed over all samples\n        count: Number of samples\n    Returns:\n        Perplexity\n\n    \"\"\"\n    return torch.exp(total / count)\n\n\ndef perplexity(preds: Tensor, target: Tensor, ignore_index: Optional[int] = None) -> Tensor:\n    \"\"\"Perplexity measures how well a language model predicts a text sample.\n\n    This metric is calculated as the average number of bits per word a model needs to represent the sample.\n\n    Args:\n        preds:\n            Logits or a unnormalized score assigned to each token in a sequence with shape [batch_size, seq_len,\n            vocab_size], which is the output of a language model. Scores will be normalized internally using softmax.\n        target:\n            Ground truth values with a shape [batch_size, seq_len].\n        ignore_index:\n            Integer specifying a target class to ignore. If given, this class index does not contribute\n            to the returned score.\n\n    Returns:\n        Perplexity value\n\n    Examples:\n        >>> import torch\n        >>> gen = torch.manual_seed(42)\n        >>> preds = torch.rand(2, 8, 5, generator=gen)\n        >>> target = torch.randint(5, (2, 8), generator=gen)\n        >>> target[0, 6:] = -100\n        >>> perplexity(preds, target, ignore_index=-100)\n        tensor(5.8540)\n\n    \"\"\"\n    total, count = _perplexity_update(preds, target, ignore_index)\n    return _perplexity_compute(total, count)\n", "path": "src/torchmetrics/functional/text/perplexity.py"}], "after_files": [{"content": "# Copyright The Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Optional, Tuple\n\nimport torch\nfrom torch import Tensor\n\n\ndef _check_shape_and_type_consistency(preds: Tensor, target: Tensor) -> None:\n    \"\"\"Check shape and type consistency of input vectors.\n\n    Args:\n        preds:\n            Logits or a unnormalized score assigned to each token in a sequence with shape [batch_size, seq_len,\n            vocab_size]. Scores will be normalized internally using softmax.\n        target:\n            Ground truth values with a shape [batch_size, seq_len].\n\n    Raises:\n        ValueError:\n            If ``preds`` tensor has no 3 dimensions.\n        ValueError:\n            If ``target`` tensor has no 2 dimensions.\n        ValueError:\n            If the first two dimensions of ``preds`` and ``target`` do not equal.\n        TypeError:\n            If ``preds`` dtype is not one of ``(torch.float16, torch.float32, torch.float64)``\n        TypeError:\n            If ``target`` is not of a type LongTensor (torch.int64)\n\n    \"\"\"\n    if len(preds.shape) != 3:\n        raise ValueError(\n            \"Input tensor `preds` is expected to have 3 dimensions, [batch_size, seq_len, vocab_size],\"\n            f\" but got {len(preds.shape)}.\"\n        )\n    if len(target.shape) != 2:\n        raise ValueError(\n            \"Input tensor `target` is expected to have 2 dimensions, [batch_size, seq_len],\"\n            f\" but got {len(target.shape)}.\"\n        )\n    if preds.shape[:2] != target.shape:\n        raise ValueError(\n            \"Input tensors `preds` and `target` are expected to have equaling first two dimensions,\"\n            f\" [batch_size, seq_len], but got {preds.shape[:2]} and {target.shape}.\"\n        )\n    if not preds.is_floating_point():\n        raise TypeError(f\"Input tensor `preds` is expected to be of floating point type but got {preds.dtype}.\")\n    if target.dtype != torch.int64:\n        raise TypeError(f\"Input tensor `target` is expected to be of a type {torch.int64} but got {target.dtype}.\")\n\n\ndef _perplexity_update(preds: Tensor, target: Tensor, ignore_index: Optional[int] = None) -> Tuple[Tensor, Tensor]:\n    \"\"\"Compute intermediate statistics for Perplexity.\n\n    Args:\n        preds:\n            Logits or a unnormalized score assigned to each token in a sequence with shape [batch_size, seq_len,\n            vocab_size]. Scores will be normalized internally using softmax.\n        target:\n            Ground truth values with a shape [batch_size, seq_len].\n        ignore_index:\n            Integer specifying a target class to ignore. If given, this class index does not contribute\n            to the returned score.\n\n    Returns:\n        Log probabilities, summed over all samples\n        Number of samples\n\n    \"\"\"\n    _check_shape_and_type_consistency(preds, target)\n\n    probs = torch.nn.functional.softmax(preds.reshape(-1, preds.shape[-1]), dim=1)\n    target = target.reshape(-1)\n\n    if ignore_index is not None:\n        mask = target.ne(ignore_index)\n        target = target.where(target != ignore_index, torch.tensor(0, device=target.device))\n    else:\n        mask = torch.ones_like(target, dtype=torch.bool)\n\n    probs = probs[torch.arange(target.numel()), target][mask]\n    total_log_probs = -probs.log().sum()\n    count = mask.sum()\n\n    return total_log_probs, count\n\n\ndef _perplexity_compute(total: Tensor, count: Tensor) -> Tensor:\n    \"\"\"Compute the Perplexity.\n\n    Args:\n        total: Log probabilities, summed over all samples\n        count: Number of samples\n    Returns:\n        Perplexity\n\n    \"\"\"\n    return torch.exp(total / count)\n\n\ndef perplexity(preds: Tensor, target: Tensor, ignore_index: Optional[int] = None) -> Tensor:\n    \"\"\"Perplexity measures how well a language model predicts a text sample.\n\n    This metric is calculated as the average number of bits per word a model needs to represent the sample.\n\n    Args:\n        preds:\n            Logits or a unnormalized score assigned to each token in a sequence with shape [batch_size, seq_len,\n            vocab_size], which is the output of a language model. Scores will be normalized internally using softmax.\n        target:\n            Ground truth values with a shape [batch_size, seq_len].\n        ignore_index:\n            Integer specifying a target class to ignore. If given, this class index does not contribute\n            to the returned score.\n\n    Returns:\n        Perplexity value\n\n    Examples:\n        >>> import torch\n        >>> gen = torch.manual_seed(42)\n        >>> preds = torch.rand(2, 8, 5, generator=gen)\n        >>> target = torch.randint(5, (2, 8), generator=gen)\n        >>> target[0, 6:] = -100\n        >>> perplexity(preds, target, ignore_index=-100)\n        tensor(5.8540)\n\n    \"\"\"\n    total, count = _perplexity_update(preds, target, ignore_index)\n    return _perplexity_compute(total, count)\n", "path": "src/torchmetrics/functional/text/perplexity.py"}]}