Repository: Lightning-AI/torchmetrics

Issue:
usage consistency on GPU between `MetricCollection` and `FeatureShare`
## 🐛 Bug

in the release notes we set them side bu side, but seems the move to GPU does not work the same, see the simple example...

### To Reproduce

```py
import torch
from torchmetrics.image import FrechetInceptionDistance, KernelInceptionDistance
from torchmetrics import MetricCollection
from torchmetrics.wrappers import FeatureShare

FrechetInceptionDistance(), KernelInceptionDistance()  # just to fetch models
img1 = torch.randint(255, (50, 3, 220, 220), dtype=torch.uint8).to("cuda")
img2 = torch.randint(255, (50, 3, 220, 220), dtype=torch.uint8).to("cuda")

mc = MetricCollection([
    FrechetInceptionDistance(),
    KernelInceptionDistance(subset_size=10, subsets=2)
]).to("cuda")
mc.update(img1, real=True)
mc.update(img2, real=False)
mc.compute()

fs = FeatureShare([
    FrechetInceptionDistance(),
    KernelInceptionDistance(subset_size=10, subsets=2)
]).to("cuda")
fs.update(img1, real=True)
fs.update(img2, real=False)
fs.compute()
```

<details>
  <summary>Error trace</summary>

```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[12], line 1
----> 1 get_ipython().run_cell_magic('timeit', '', 'from torchmetrics.wrappers import FeatureShare\n\nfs = FeatureShare([\n    FrechetInceptionDistance(),\n    KernelInceptionDistance(subset_size=10, subsets=2)\n]).to("cuda")\nfs.update(img1, real=True)\nfs.update(img2, real=False)\nfs.compute()\n')

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/IPython/core/interactiveshell.py:2515, in InteractiveShell.run_cell_magic(self, magic_name, line, cell)
   2513 with self.builtin_trap:
   2514     args = (magic_arg_s, cell)
-> 2515     result = fn(*args, **kwargs)
   2517 # The code below prevents the output from being displayed
   2518 # when using magics with decorator @output_can_be_silenced
   2519 # when the last Python token in the expression is a ';'.
   2520 if getattr(fn, magic.MAGIC_OUTPUT_CAN_BE_SILENCED, False):

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/IPython/core/magics/execution.py:1185, in ExecutionMagics.timeit(self, line, cell, local_ns)
   1183 for index in range(0, 10):
   1184     number = 10 ** index
-> 1185     time_number = timer.timeit(number)
   1186     if time_number >= 0.2:
   1187         break

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/IPython/core/magics/execution.py:173, in Timer.timeit(self, number)
    171 gc.disable()
    172 try:
--> 173     timing = self.inner(it, self.timer)
    174 finally:
    175     if gcold:

File <magic-timeit>:7, in inner(_it, _timer)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torchmetrics/collections.py:220, in MetricCollection.update(self, *args, **kwargs)
    218 for m in self.values(copy_state=False):
    219     m_kwargs = m._filter_kwargs(**kwargs)
--> 220     m.update(*args, **m_kwargs)
    222 if self._enable_compute_groups:
    223     self._merge_compute_groups()

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torchmetrics/metric.py:475, in Metric._wrap_update.<locals>.wrapped_func(*args, **kwargs)
    467         if "Expected all tensors to be on" in str(err):
    468             raise RuntimeError(
    469                 "Encountered different devices in metric calculation (see stacktrace for details)."
    470                 " This could be due to the metric class not being on the same device as input."
   (...)
    473                 " device corresponds to the device of the input."
    474             ) from err
--> 475         raise err
    477 if self.compute_on_cpu:
    478     self._move_list_states_to_cpu()

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torchmetrics/metric.py:465, in Metric._wrap_update.<locals>.wrapped_func(*args, **kwargs)
    463 with torch.set_grad_enabled(self._enable_grad):
    464     try:
--> 465         update(*args, **kwargs)
    466     except RuntimeError as err:
    467         if "Expected all tensors to be on" in str(err):

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torchmetrics/image/fid.py:335, in FrechetInceptionDistance.update(self, imgs, real)
    333 """Update the state with extracted features."""
    334 imgs = (imgs * 255).byte() if self.normalize else imgs
--> 335 features = self.inception(imgs)
    336 self.orig_dtype = features.dtype
    337 features = features.double()

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1525         or _global_backward_pre_hooks or _global_backward_hooks
   1526         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527     return forward_call(*args, **kwargs)
   1529 try:
   1530     result = None

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torchmetrics/wrappers/feature_share.py:41, in NetworkCache.forward(self, *args, **kwargs)
     39 def forward(self, *args: Any, **kwargs: Any) -> Any:
     40     """Call the network with the given arguments."""
---> 41     return self.network(*args, **kwargs)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1525         or _global_backward_pre_hooks or _global_backward_hooks
   1526         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527     return forward_call(*args, **kwargs)
   1529 try:
   1530     result = None

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torchmetrics/image/fid.py:155, in NoTrainInceptionV3.forward(self, x)
    153 def forward(self, x: Tensor) -> Tensor:
    154     """Forward pass of neural network with reshaping of output."""
--> 155     out = self._torch_fidelity_forward(x)
    156     return out[0].reshape(x.shape[0], -1)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torchmetrics/image/fid.py:90, in NoTrainInceptionV3._torch_fidelity_forward(self, x)
     83 x = interpolate_bilinear_2d_like_tensorflow1x(
     84     x,
     85     size=(self.INPUT_IMAGE_SIZE, self.INPUT_IMAGE_SIZE),
     86     align_corners=False,
     87 )
     88 x = (x - 128) / 128
---> 90 x = self.Conv2d_1a_3x3(x)
     91 x = self.Conv2d_2a_3x3(x)
     92 x = self.Conv2d_2b_3x3(x)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1525         or _global_backward_pre_hooks or _global_backward_hooks
   1526         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527     return forward_call(*args, **kwargs)
   1529 try:
   1530     result = None

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch_fidelity/feature_extractor_inceptionv3.py:208, in BasicConv2d.forward(self, x)
    207 def forward(self, x):
--> 208     x = self.conv(x)
    209     x = self.bn(x)
    210     return F.relu(x, inplace=True)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1525         or _global_backward_pre_hooks or _global_backward_hooks
   1526         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527     return forward_call(*args, **kwargs)
   1529 try:
   1530     result = None

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/conv.py:460, in Conv2d.forward(self, input)
    459 def forward(self, input: Tensor) -> Tensor:
--> 460     return self._conv_forward(input, self.weight, self.bias)

File /home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/nn/modules/conv.py:456, in Conv2d._conv_forward(self, input, weight, bias)
    452 if self.padding_mode != 'zeros':
    453     return F.conv2d(F.pad(input, self._reversed_padding_repeated_twice, mode=self.padding_mode),
    454                     weight, bias, self.stride,
    455                     _pair(0), self.dilation, self.groups)
--> 456 return F.conv2d(input, weight, bias, self.stride,
    457                 self.padding, self.dilation, self.groups)

RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same
```

</details>

### Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

### Environment

- TorchMetrics version: `1.3.0.post`
- Python & PyTorch Version (e.g., 1.0): Python 3.10.10
- Any other relevant information such as OS (e.g., Linux): linux

### Additional context

<!-- Add any other context about the problem here. -->



Discussion:
Hi! thanks for your contribution!, great first issue!

Resolving patch:
diff --git a/src/torchmetrics/wrappers/feature_share.py b/src/torchmetrics/wrappers/feature_share.py
--- a/src/torchmetrics/wrappers/feature_share.py
+++ b/src/torchmetrics/wrappers/feature_share.py
@@ -34,7 +34,8 @@ class NetworkCache(Module):
     def __init__(self, network: Module, max_size: int = 100) -> None:
         super().__init__()
         self.max_size = max_size
-        self.network = lru_cache(maxsize=self.max_size)(network)
+        self.network = network
+        self.network.forward = lru_cache(maxsize=self.max_size)(network.forward)
 
     def forward(self, *args: Any, **kwargs: Any) -> Any:
         """Call the network with the given arguments."""


Test patch:
diff --git a/tests/unittests/wrappers/test_feature_share.py b/tests/unittests/wrappers/test_feature_share.py
--- a/tests/unittests/wrappers/test_feature_share.py
+++ b/tests/unittests/wrappers/test_feature_share.py
@@ -95,14 +95,14 @@ def test_memory():
 
     fid = FrechetInceptionDistance(feature=64).cuda()
     inception = InceptionScore(feature=64).cuda()
-    kid = KernelInceptionDistance(feature=64).cuda()
+    kid = KernelInceptionDistance(feature=64, subset_size=5).cuda()
 
     memory_before_fs = torch.cuda.memory_allocated()
     assert memory_before_fs > base_memory, "The memory usage should be higher after initializing the metrics."
 
     torch.cuda.empty_cache()
 
-    FeatureShare([fid, inception, kid]).cuda()
+    feature_share = FeatureShare([fid, inception, kid]).cuda()
     memory_after_fs = torch.cuda.memory_allocated()
 
     assert (
@@ -112,6 +112,19 @@ def test_memory():
         memory_after_fs < memory_before_fs
     ), "The memory usage should be higher after initializing the feature share wrapper."
 
+    img1 = torch.randint(255, (50, 3, 220, 220), dtype=torch.uint8).to("cuda")
+    img2 = torch.randint(255, (50, 3, 220, 220), dtype=torch.uint8).to("cuda")
+
+    feature_share.update(img1, real=True)
+    feature_share.update(img2, real=False)
+    res = feature_share.compute()
+
+    assert "cuda" in str(res["FrechetInceptionDistance"].device)
+    assert "cuda" in str(res["InceptionScore"][0].device)
+    assert "cuda" in str(res["InceptionScore"][1].device)
+    assert "cuda" in str(res["KernelInceptionDistance"][0].device)
+    assert "cuda" in str(res["KernelInceptionDistance"][1].device)
+
 
 def test_same_result_as_individual():
     """Test that the feature share wrapper gives the same result as the individual metrics."""