code: def test_auc_reset(self) -> None:
      batch_size = 64
      auc = AUCMetric(
          world_size=1,
          my_rank=0,
          batch_size=batch_size,
          tasks=[DefaultTaskInfo],
      )
      # Mimic the case where the metric module is moved to GPU.
      device = torch.device("cuda:0")
      auc.to(device)
      # Mimic that reset is called when checkpointing.
      auc.reset()
      self.assertEqual(len(auc._metrics_computations[0].predictions), 1)
      self.assertEqual(len(auc._metrics_computations[0].labels), 1)
      self.assertEqual(len(auc._metrics_computations[0].weights), 1)
      model_output = gen_test_batch(batch_size)
      model_output = {k: v.to(device) for k, v in model_output.items()}
      auc.update(
          predictions={"DefaultTask": model_output["prediction"]},
          labels={"DefaultTask": model_output["label"]},
          weights={"DefaultTask": model_output["weight"]},
      )
      self.assertEqual(len(auc._metrics_computations[0].predictions), 1)
      self.assertEqual(len(auc._metrics_computations[0].labels), 1)
      self.assertEqual(len(auc._metrics_computations[0].weights), 1)
      self.assertEqual(
          auc._metrics_computations[0].predictions[0].device, device
      )
      self.assertEqual(
          auc._metrics_computations[0].labels[0].device, device
      )
      self.assertEqual(
          auc._metrics_computations[0].weights[0].device, device
      )
  
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 1 for calling the TestGPU obj's underlying member method code with example usage: obj.test_auc_reset() without return types