def checkpoint_wrapper(m):
    """
    A friendlier wrapper for performing activation checkpointing.

    Compared to the PyTorch version, this version:
    - wraps an nn.Module, so that all subsequent calls will use checkpointing
    - handles keyword arguments in the forward
    - handles non-Tensor outputs from the forward

    Usage::

        checkpointed_module = checkpoint_wrapper(my_module)
        a, b = checkpointed_module(x, y=3, z=torch.Tensor([1]))
    """
    original_forward = m.forward

    def _checkpointed_forward(*args, **kwargs):
        # Autograd Functions in PyTorch work best with positional args, since
        # the backward must return gradients (or None) for every input argument.
        # We can flatten keyword arguments to make this easier.
        kwarg_keys, flat_args = pack_kwargs(*args, **kwargs)
        parent_ctx_dict = {}
        output = CheckpointFunction.apply(
            original_forward, parent_ctx_dict, kwarg_keys, *flat_args
        )
        if isinstance(output, torch.Tensor):
            return output
        else:
            packed_non_tensor_outputs = parent_ctx_dict["packed_non_tensor_outputs"]
            if packed_non_tensor_outputs:
                output = unpack_non_tensors(output, packed_non_tensor_outputs)
            return output

    m.forward = _checkpointed_forward
    return m

    A friendlier wrapper for performing activation checkpointing.

    Compared to the PyTorch version, this version:
    - wraps an nn.Module, so that all subsequent calls will use checkpointing
    - handles keyword arguments in the forward
    - handles non-Tensor outputs from the forward

    Usage::

        checkpointed_module = checkpoint_wrapper(my_module)
        a, b = checkpointed_module(x, y=3, z=torch.Tensor([1]))
    
decoding/IAD/fairseq/fairseq/modules/checkpoint_activations.py