Update test_label_tensor.py

tests/test_label_tensor.py

@@ -27,7 +27,6 @@ def test_labels():
 def test_extract():
     label_to_extract = ['a', 'c']
     tensor = LabelTensor(data, labels)
-    print(tensor)
     new = tensor.extract(label_to_extract)
     assert new.labels == label_to_extract
     assert new.shape[1] == len(label_to_extract)
@@ -58,7 +57,6 @@ def test_extract_order():
     expected = torch.cat(
         (data[:, 2].reshape(-1, 1), data[:, 0].reshape(-1, 1)),
         dim=1)
-    print(expected)
     assert new.labels == label_to_extract
     assert new.shape[1] == len(label_to_extract)
     assert torch.all(torch.isclose(expected, new))
@@ -83,6 +81,13 @@ def test_merge2():
 
 
 def test_getitem():
+    tensor = LabelTensor(data, labels)
+    tensor_view = tensor['a']
+
+    assert tensor_view.labels == ['a']
+    assert torch.allclose(tensor_view.flatten(), data[:, 0])
+
+def test_getitem2():
     tensor = LabelTensor(data, labels)
     tensor_view = tensor[:5]
 
@@ -101,4 +106,4 @@ def test_slice():
 
     tensor_view3 = tensor[:, 2]
     assert tensor_view3.labels == labels[2]
-    assert torch.allclose(tensor_view3, data[:, 2].reshape(-1, 1))
+    assert torch.allclose(tensor_view3, data[:, 2].reshape(-1, 1))