[minor] test split_non_tensor(<Tensor>) case (#1000)

facebookresearch · May 30, 2022 · e7602a4 · e7602a4
1 parent 2bd85c0
commit e7602a4
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/fairscale/utils/containers.py b/fairscale/utils/containers.py
@@ -78,6 +78,8 @@ def split_non_tensors(
     Split a tuple into a list of tensors and the rest with information
     for later reconstruction.
 
+    When called with a tensor X, will return: (x,), None
+
     Usage::
 
         x = torch.Tensor([1])

diff --git a/tests/utils/test_containers.py b/tests/utils/test_containers.py
@@ -99,6 +99,11 @@ def test_split_unpack():
     x = torch.Tensor([1])
     y = torch.Tensor([2])
 
+    # degenerate case, args is a single tensor.
+    tensors, packed_non_tensors = split_non_tensors(x)
+    assert tensors == (x,)
+    assert packed_non_tensors is None
+
     tensors, packed_non_tensors = split_non_tensors((x, y, None, 3))
     assert tensors == (x, y)
     assert packed_non_tensors == {