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_current_stream(self):
      # Test current stream on the device and check if the stream device index
      # matches with the device ID
      @torch.jit.script
      def fn():
          device_index = torch.cuda.current_device()
          device = torch.device("cuda:" + str(device_index))
          s0 = torch.cuda.current_stream(device)
          s1 = torch.cuda.current_stream(torch.device("cuda:1"))
          s2 = torch.cuda.current_stream(torch.device("cuda:0"))
  
          return s0.device_index(), s1.device_index(), s2.device_index()
  
      d0, d1, d2 = fn()
      # By default, the current device ID is 0.
      self.assertEqual(0, d0)
      self.assertEqual(1, d1)
      self.assertEqual(0, d2)
      self.assertEqual(d0, d2)
  
      # Test current_stream API by passing device ID as an argument and
      # and check if the stream device index matches with the device ID
      @torch.jit.script
      def fn_with_device_index_args():
          device_index = torch.cuda.current_device()
          s0 = torch.cuda.current_stream(device_index)
          s1 = torch.cuda.current_stream(1)
          s2 = torch.cuda.current_stream(0)
  
          return s0.device_index(), s1.device_index(), s2.device_index()
  
      d0, d1, d2 = fn_with_device_index_args()
      # By default, the current device ID is 0.
      self.assertEqual(0, d0)
      self.assertEqual(1, d1)
      self.assertEqual(0, d2)
      self.assertEqual(d0, d2)
  
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 5 for calling the TestCUDA obj's underlying member method code with example usage: obj.test_current_stream() and returns: s0,  s1,  s2, s0,  s1,  s2