def test_class_filtering(self, sample_boxes):
        """Test filtering based on class IDs."""
        # Set classes_to_keep to [1, 2]
        filter = BoxFiltering(classes_to_keep=tensor([1, 2]))
        filtered_boxes = filter(sample_boxes)

        # Expected output: only boxes with class_id 1 and 2 should be kept
        assert len(filtered_boxes[0]) == 2  # Two boxes in the first batch
        assert filtered_boxes[0][0][0].item() == 1  # Box with class_id 1
        assert filtered_boxes[0][1][0].item() == 2  # Box with class_id 2
        assert len(filtered_boxes[1]) == 2  # Two boxes in the second batch
        assert filtered_boxes[1][0][0].item() == 1  # Box with class_id 1
        assert filtered_boxes[1][1][0].item() == 2  # Box with class_id 2
        assert len(filtered_boxes[2]) == 2  # Two boxes in the third batch
        assert filtered_boxes[2][0][0].item() == 1  # Box with class_id 1
        assert filtered_boxes[2][1][0].item() == 2  # Box with class_id 2
Test filtering based on class IDs.
tests/models/box_filtering.py