    def wrap_module(self, module: nn.Module, **kwargs) -> nn.Module:
        """
        Wrap a module so that the forward method uses PyTorch's [checkpointing functionality]
        (https://pytorch.org/docs/stable/checkpoint.html).

        !!! Note
            Currently this `CheckpointWrapper` implementation requires that the wrapped
            module is called with positional arguments only.

            We recommend you use the
            :class:`allennlp.nn.checkpoint.fairscale_checkpoint_wrapper.FairScaleCheckpointWrapper`
            if you need more flexibility.

        """
        # Inspired by FairScale:
        #  --> https://github.com/facebookresearch/fairscale/blob/1e4a503cda8571851a68effd6e504a192838ab06/fairscale/nn/checkpoint/checkpoint_activations.py#L145-L153  # noqa: E501
        # We just patch the forward method to avoid having to proxy all the fields and other methods.
        # The use of weakref here is to prevent creating a ref cycle: m -> m.forward -> m.

        assert len(kwargs) == 0  # This way of wrapping only works for positional arguments.

        module.forward = functools.partial(  # type: ignore[assignment]
            _checkpointed_forward, type(module).forward, weakref.ref(module)
        )
        return module
Wrap a module so that the forward method uses PyTorch's [checkpointing functionality]
(https://pytorch.org/docs/stable/checkpoint.html).

!!! Note
Currently this `CheckpointWrapper` implementation requires that the wrapped
module is called with positional arguments only.

We recommend you use the
:class:`allennlp.nn.checkpoint.fairscale_checkpoint_wrapper.FairScaleCheckpointWrapper`
if you need more flexibility.
allennlp/nn/checkpoint/checkpoint_wrapper.py