s1 = torch.cuda.Stream(device=1)
s2 = torch.cuda.Stream(device=0)
+ with torch.cuda.device(s1.device):
+ prev_stream_on_cuda1 = torch.cuda.current_stream()
+
self.assertEqual(torch.cuda.current_stream(), s0)
self.assertEqual(0, torch.cuda.current_device())
with torch.cuda.stream(s1):
self.assertEqual(torch.cuda.current_stream(), s1)
self.assertEqual(1, torch.cuda.current_device())
+ with torch.cuda.device(s1.device):
+ self.assertEqual(prev_stream_on_cuda1, torch.cuda.current_stream())
+
self.assertEqual(torch.cuda.current_stream(), s0)
self.assertEqual(0, torch.cuda.current_device())
if stream is None:
yield
return
- prev_stream = current_stream()
+ src_prev_stream = current_stream()
+
+ if src_prev_stream.device != stream.device:
+ # The given stream is on a different device; have to restore the
+ # current_stream on that device on exit as well
+ with device(stream.device):
+ dst_prev_stream = current_stream()
+
torch._C._cuda_setStream(stream._cdata)
try:
yield
finally:
- torch._C._cuda_setStream(prev_stream._cdata)
+ if src_prev_stream.device != stream.device:
+ torch._C._cuda_setStream(dst_prev_stream._cdata)
+ torch._C._cuda_setStream(src_prev_stream._cdata)
def device_count():