def test_count(self):
        """Test that count 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(da.count, {0: 10, 1: 10})
Test that count is correctly computed.
tests/benchmarks/utils/test_data_attribute.py