code: def test_cpu_and_ranks_only(self):
      device = torch.device("cuda")
      state_dict = {
          "tensor1": torch.arange(10, device=device),
          "tensor2": torch.ones(10, device=device),
      }
  
      cpu_state_dict = _offload_state_dict_to_cpu(
          state_dict, ranks_only=(0, 2)
      )
      if dist.get_rank() in (0, 2):
          for v in cpu_state_dict.values():
              self.assertFalse(v.is_cuda)
          self.assertEqual(cpu_state_dict["tensor1"], torch.arange(10))
          self.assertEqual(cpu_state_dict["tensor2"], torch.ones(10))
      else:
          self.assertEqual(cpu_state_dict, {})
  
file_path: pytorch/test/distributed/checkpoint/test_state_dict_utils.py
imports: import torch
  import torch.distributed
  import torch.distributed._functional_collectives
  from torch.distributed._state_dict_utils import _gather_state_dict, _offload_state_dict_to_cpu
  from torch.distributed._tensor import DTensor
  from torch.distributed._tensor.placement_types import Shard
  from torch.testing._internal.common_utils import run_tests
  from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase, skip_if_lt_x_gpu, with_comms
label_desc: Use image node_id 4 for calling the TestStateDictUtils obj's underlying member method code with example usage: obj.test_cpu_and_ranks_only() without return types