diff --git a/notebooks/text-generation/scripts/run_clm.py b/notebooks/text-generation/scripts/run_clm.py
index 03ff71f80..5e2c6744a 100644
--- a/notebooks/text-generation/scripts/run_clm.py
+++ b/notebooks/text-generation/scripts/run_clm.py
@@ -1,14 +1,13 @@
 from dataclasses import dataclass, field
-import os
+
+from datasets import load_from_disk
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
-    set_seed,
     default_data_collator,
+    set_seed,
 )
-from datasets import load_from_disk
-import torch
-from huggingface_hub import login
+
 from optimum.neuron import NeuronHfArgumentParser as HfArgumentParser
 from optimum.neuron import NeuronTrainer as Trainer
 from optimum.neuron import NeuronTrainingArguments as TrainingArguments
diff --git a/notebooks/text-generation/scripts/utils/pack_dataset.py b/notebooks/text-generation/scripts/utils/pack_dataset.py
index 38220389a..9f06a8637 100644
--- a/notebooks/text-generation/scripts/utils/pack_dataset.py
+++ b/notebooks/text-generation/scripts/utils/pack_dataset.py
@@ -1,5 +1,6 @@
-from itertools import chain
 from functools import partial
+from itertools import chain
+
 
 remainder = {"input_ids": [], "attention_mask": [], "token_type_ids": []}
 
diff --git a/optimum/neuron/modeling_decoder.py b/optimum/neuron/modeling_decoder.py
index e9e3c495d..d54274a39 100644
--- a/optimum/neuron/modeling_decoder.py
+++ b/optimum/neuron/modeling_decoder.py
@@ -246,7 +246,7 @@ def _save_pretrained(self, save_directory: Union[str, Path]):
 
         if src_compiled_path is None:
             # The compiled model has never been serialized: do it now
-            self.model._save_compiled_artifacts(dst_compiled_path)
+            self.model.save(dst_compiled_path)
         else:
             shutil.copytree(src_compiled_path, dst_compiled_path)