def test_subset(self):
        """Test that subset is correctly computed."""
        t0 = torch.zeros(10, dtype=torch.int)
        t1 = torch.ones(10, dtype=torch.int)
        da = DataAttribute(torch.cat([t0, t1]), "task_labels")
        self.assertEqual(list(da.subset(range(10)).data), list(t0))
        self.assertEqual(list(da.subset(range(10, 20)).data), list(t1))
Test that subset is correctly computed.
tests/benchmarks/utils/test_data_attribute.py