Python default_stream Examples

Example #1

File: test_portal.py Project: yuhc/ava-pytorch

def test_copy_returns_on_next_device():
    portal = Portal(torch.rand(1), tensor_life=1)

    prev_stream = default_stream(torch.device("cpu"))
    next_stream = default_stream(torch.device("cuda"))

    phony = torch.zeros(0, requires_grad=True)
    assert phony.device.type == "cpu"

    phony = portal.copy(prev_stream, next_stream, phony)
    assert phony.device.type == "cuda"

Example #2

File: test_stream.py Project: zacker150/pytorch

 def test_default_stream_cuda(self):
     stream = default_stream(torch.device("cuda"))
     assert isinstance(stream, torch.cuda.Stream)
     assert stream == torch.cuda.default_stream()

Example #3

File: test_stream.py Project: zacker150/pytorch

 def test_default_stream_cpu(self):
     stream = default_stream(torch.device("cpu"))
     assert stream is CPUStream