adding test

mathLab · Oct 31, 2023 · 84aaf7b · 84aaf7b
1 parent 8bff3e9
commit 84aaf7b
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tests/test_label_tensor.py b/tests/test_label_tensor.py
@@ -87,6 +87,11 @@ def test_getitem():
     assert tensor_view.labels == ['a']
     assert torch.allclose(tensor_view.flatten(), data[:, 0])
 
+    tensor_view = tensor['a', 'c']
+
+    assert tensor_view.labels == ['a', 'c']
+    assert torch.allclose(tensor_view, data[:, 0::2])
+
 def test_getitem2():
     tensor = LabelTensor(data, labels)
     tensor_view = tensor[:5]