def model(self):
        """Only load the model if needed."""
        model = PLBartForConditionalGeneration.from_pretrained(self.checkpoint_name).to(torch_device)
        if "cuda" in torch_device:
            model = model.half()
        return model
Only load the model if needed.
dpkd/transformers/tests/models/plbart/test_modeling_plbart.py