def load_model(self, file_name):
        """
        Loads a PyTorch module's state.

        Args:
            file_name (str): File name of saved the model's `state_dict()`
        """

        model_to_load = (
            self.model.module if hasattr(self.model, "module") else self.model
        )  # Take care of distributed/parallel training

        model_to_load.load_state_dict(torch.load(file_name))
        logger.info("Model checkpoint loaded from %s", file_name)

        Loads a PyTorch module's state.

        Args:
            file_name (str): File name of saved the model's `state_dict()`
        
utils_nlp/models/transformers/common.py