def test_filter_as_zero(self, sample_boxes):
        """Test filtering boxes as zero when filter_as_zero is True."""
        filter = BoxFiltering(confidence_threshold=0.8, filter_as_zero=True)
        filtered_boxes = filter(sample_boxes)

        # Expected output: boxes with confidence <= 0.8 should be zeroed out
        assert torch.all(filtered_boxes[0][0] != 0)  # Box with confidence 0.9 should remain
        assert torch.all(filtered_boxes[0][1:] == 0)  # Remaining boxes should be zeroed
        assert torch.all(filtered_boxes[1][0] != 0)  # Box with confidence 0.95 should remain
        assert torch.all(filtered_boxes[1][1:] == 0)  # Remaining boxes should be zeroed
        assert torch.all(filtered_boxes[2] == 0)  # All boxes in the third batch should be zeroed
Test filtering boxes as zero when filter_as_zero is True.
tests/models/box_filtering.py