    def resize_token_embeddings(
        self, new_size: int, *, init_fn: Callable = torch.nn.init.normal_
    ) -> None:
        """
        Resizes the token embeddings in the model.

        This takes care of the token embeddings for the encoder, the decoder, and the LM head.

        new_size : `int`
            The new size of the token embeddings
        init_fn : `Callable`
            The function to use to initialize new embeddings. This function will be called with a
            single argument, the tensor to initialize, and it is expected to initialize the tensor
            in place. Many of the functions from `torch.nn.init` fit.
        """
        self.encoder.resize_token_embeddings(new_size, init_fn=init_fn)
        # If encoder and decoder share embeddings, this is a no-op the second time.
        self.decoder.resize_token_embeddings(new_size, init_fn=init_fn)

        # resize lm head
        old_size = self.lm_head.out_features
        if old_size == new_size:
            return
        new_lm_head = torch.nn.Linear(
            self.lm_head.in_features,
            new_size,
            self.lm_head.bias,
            self.lm_head.weight.device,
            self.lm_head.weight.dtype,
        )
        copy_size = min(old_size, new_size)
        new_lm_head.weight.data[:copy_size, ...] = self.lm_head.weight.data[:copy_size, ...]
        if self.lm_head.bias and new_lm_head.bias:
            new_lm_head.bias.data[:copy_size, ...] = self.lm_head.bias[:copy_size, ...]
        if new_size > old_size:
            init_fn(new_lm_head.weight.data[copy_size:, ...])
            if new_lm_head.bias:
                init_fn(new_lm_head.bias[copy_size:, ...])

        self.lm_head = new_lm_head
Resizes the token embeddings in the model.

This takes care of the token embeddings for the encoder, the decoder, and the LM head.

new_size : `int`
The new size of the token embeddings
init_fn : `Callable`
The function to use to initialize new embeddings. This function will be called with a
single argument, the tensor to initialize, and it is expected to initialize the tensor
in place. Many of the functions from `torch.nn.init` fit.
allennlp/modules/transformer/t5.py