Ejemplo n.º 1
0
def test_copy_wait_cuda_cuda(cuda_sleep):
    prev_stream = current_stream(torch.device("cuda"))
    next_stream = new_stream(torch.device("cuda"))
    _test_copy_wait(prev_stream, next_stream, cuda_sleep)
Ejemplo n.º 2
0
 def test_get_device_cuda(self):
     stream = current_stream(torch.device("cuda"))
     assert get_device(stream).type == "cuda"
Ejemplo n.º 3
0
 def test_use_stream_cuda(self):
     stream = new_stream(torch.device("cuda"))
     with use_stream(stream):
         assert current_stream(torch.device("cuda")) == stream
Ejemplo n.º 4
0
 def test_current_stream_cuda(self):
     stream = current_stream(torch.device("cuda"))
     assert isinstance(stream, torch.cuda.Stream)
     assert stream == torch.cuda.current_stream()
Ejemplo n.º 5
0
 def test_current_stream_cpu(self):
     stream = current_stream(torch.device("cpu"))
     assert stream is CPUStream
Ejemplo n.º 6
0
 def test_wait_stream_cuda_cuda(self, cuda_sleep):
     source = current_stream(torch.device("cuda"))
     target = new_stream(torch.device("cuda"))
     self._test_wait_stream(source, target, cuda_sleep)