code: def test_apply_to_tensors(devices):
      """Test apply_to_tensors for both cpu & gpu"""
      if (
          "cuda" in devices
          and not torch.cuda.is_available()
          or torch.cuda.device_count() < 1
      ):
          pytest.skip("Skipped due to lack of GPU")
      expected = 0
  
      def get_a_tensor():
          """Return a random tensor on random device."""
          dev = random.choice(devices)
          shape = random.choice((1, (2, 3), (4, 5, 6), (7, 8, 9, 10)))
          t = torch.rand(shape).to(dev)
          nonlocal expected
          expected += t.numel()
          return t
  
      # create a mixed bag of data.
      data = [1, "str"]  # list
      # dict
      data.append(
          {
              "key1": get_a_tensor(),
              "key2": {1: get_a_tensor()},
              "key3": 3,
          }
      )
      # set
      data.insert(0, set(["x", get_a_tensor(), get_a_tensor()]))
      # tuple
      data.append(
          ([1], get_a_tensor(), 1, [get_a_tensor()], set((1, 2)))
      )
      # OrderedDict
      od = OrderedDict()
      od["k"] = "value"
      data.append(od)
      # namedtuple
      NT = namedtuple("NT", ["key1", "key2"])
      nt = NT(key1=1, key2=get_a_tensor())
      data.append(nt)
  
      total = 0
  
      def fn(t, x=[[total]]):
          nonlocal total
          total += t.numel()
          return t
  
      new_data = apply_to_tensors(fn, data)
      assert total == expected, f"{total} vs. {expected}"
      for i, v in enumerate(data):
          assert type(new_data[i]) == type(
              v
          ), f"expected type {type(v)} got {type(new_data[i])}"
  
file_path: fairscale/tests/utils/test_containers.py
imports: from collections import OrderedDict, namedtuple
  import random
  import pytest
  import torch
  import torch.nn
  from fairscale.internal.containers import apply_to_tensors, pack_kwargs, split_non_tensors, unpack_kwargs, unpack_non_tensors
label_desc: Use image node_id 1 for calling a global function with example usage: test_apply_to_tensors(devices) and returns: t, t