def METHOD_NAME(self, state_dict, name):
    super().METHOD_NAME(state_dict, name)
    prefix = name + "." if name != "" else ""
    # some old checkpoints had weight sharing implemented incorrectly
    # (note: this was correct in the original paper code)
    if utils.item(state_dict.get(f"{prefix}version", torch.tensor(1))) < 2:
        state_dict[f"{prefix}version"] = torch.tensor(1)
        # check if input embeddings and output embeddings were tied
        if not torch.allclose(
            state_dict[f"{prefix}sentence_encoder.embed_tokens.weight"],
            state_dict[f"{prefix}lm_head.weight"],
        ):
            # they weren't tied, re-init the LM head without weight sharing
            self.lm_head = self.build_lm_head(
                embed_dim=self.args.encoder_embed_dim,
                output_dim=len(self.dictionary),
                activation_fn=self.args.activation_fn,
                weight=None,  # don't share weights
            )