def token_dropper(cls, x, mask):
                """
                The implementation of dropping tokens has been done slightly differently in this tool.
                Here we check that it is equivalent to the original implementation
                """

                class MaskProcessor(torch.nn.Module):
                    def __init__(self, patch_size=16):
                        super().__init__()
                        self.avg_pool = torch.nn.AvgPool2d(patch_size)

                    def forward(self, ones_mask):
                        B = ones_mask.shape[0]
                        ones_mask = ones_mask[0].unsqueeze(0)  # take the first mask
                        ones_mask = self.avg_pool(ones_mask)[0]
                        ones_mask = torch.where(ones_mask.view(-1) > 0)[0] + 1
                        ones_mask = torch.cat([torch.IntTensor(1).fill_(0).to(device), ones_mask]).unsqueeze(0)
                        ones_mask = ones_mask.expand(B, -1)
                        return ones_mask

                mask_processor = MaskProcessor()
                patch_mask = mask_processor(mask)

                # x = self.pos_drop(x) # B, N, C
                if patch_mask is not None:
                    # patch_mask is B, K
                    B, N, C = x.shape
                    if len(patch_mask.shape) == 1:  # not a separate one per batch
                        x = x[:, patch_mask]
                    else:
                        patch_mask = patch_mask.unsqueeze(-1).expand(-1, -1, C)
                        x = torch.gather(x, 1, patch_mask)
                return x

                The implementation of dropping tokens has been done slightly differently in this tool.
                Here we check that it is equivalent to the original implementation
                
tests/estimators/certification/test_vision_transformers.py