class_docstr: A suite of tests for the CUDA API in TorchScript.
class_docstr_tok: ["A","suite","of","tests","for","the","CUDA","API","in","TorchScript","."]
code: def test__maybe_exchange_device_op(self):
      def fn(device: int, tensor):
          torch.cuda._maybe_exchange_device(device)
          return tensor.cos().relu()
  
      fn_s = torch.jit.script(fn)
      # Just check the graph, don't run it. Otherwise, we'd  need to
      # run this test on a multi-gpu CI runner, which is overkill.
      g = fn_s.graph
      FileCheck().check("cuda::_maybe_exchange_device(").run(g)
      torch._C._jit_pass_inline(g)
      FileCheck().check("cuda::_maybe_exchange_device(").run(g)
  
file_path: pytorch/test/jit/test_cuda.py
imports: import os
  import sys
  import gc
  import unittest
  import torch
  from typing import NamedTuple
  from torch.testing import FileCheck
  from torch.testing._internal.jit_utils import JitTestCase
  from torch.testing._internal.common_utils import skipIfRocm, skipCUDANonDefaultStreamIf, NoTest, TEST_CUDA
  from torch.testing._internal.common_cuda import TEST_MULTIGPU
label_desc: Use image node_id 8 for calling the TestCUDA obj's underlying member method code with example usage: obj.test__maybe_exchange_device_op() and returns: tensor