def METHOD_NAME(self, mask):
    module = DropMaskLinear(HIDDEN_DIM, HIDDEN_DIM)
    inputs = (torch.rand(HIDDEN_DIM), None)
    inputs[0].requires_grad = True
    _test_activation_checkpoint(module, *inputs)