code: def test_device_match(self) -> None:
      config = copy.deepcopy(DefaultMetricsConfig)
      config.rec_metrics[RecMetricEnum.NE].window_size = 256
      batch_size = 128
      device = torch.device("cuda:0")
      metric_module = generate_metric_module(
          RecMetricModule,
          metrics_config=config,
          batch_size=batch_size,
          world_size=1,
          my_rank=0,
          state_metrics_mapping={},
          device=device,
      )
  
      for _ in range(10):
          model_output = gen_test_batch(batch_size)
          model_output = {
              k: v.to(device) for k, v in model_output.items()
          }
          metric_module.update(model_output)
  
file_path: torchrec/torchrec/metrics/tests/test_gpu.py
imports: import copy
  import unittest
  import torch
  from torchrec.metrics.auc import AUCMetric
  from torchrec.metrics.metric_module import generate_metric_module, RecMetricModule
  from torchrec.metrics.metrics_config import DefaultMetricsConfig, DefaultTaskInfo, RecMetricEnum
  from torchrec.metrics.ne import NEMetric
  from torchrec.metrics.test_utils import gen_test_batch
label_desc: Use image node_id 2 for calling the TestGPU obj's underlying member method code with example usage: obj.test_device_match() without return types