Test default CUDA index detection of gather/scatter

tests/test_microbatch.py

@@ -1,4 +1,6 @@
+import pytest
 import torch
+import torch.cuda
 
 from torchgpipe.microbatch import gather, scatter
 
@@ -43,3 +45,17 @@ def test_scatter_tuple():
     assert b[0].size() == (1, 1)
     assert a[1].size() == (2, 2)
     assert b[1].size() == (2, 2)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason='cuda required')
+def test_default_device_index():
+    default_cuda = torch.device('cuda')
+    assert default_cuda.index is None
+
+    x = torch.rand(2, 1)
+    a, b = scatter(x, chunks=2, device=default_cuda)
+    y = gather([a, b], device=default_cuda)
+
+    assert a.is_cuda
+    assert b.is_cuda
+    assert y.is_cuda