diff --git a/.gitignore b/.gitignore
index eeb1def..f698b5a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -11,4 +11,4 @@ wandb/
 *.pyc
 *.pkl
 .vscode/
-checkpoints/
\ No newline at end of file
+checkpoints/
diff --git a/train_accelerate.py b/train_accelerate.py
index 0531d61..7279114 100644
--- a/train_accelerate.py
+++ b/train_accelerate.py
@@ -24,7 +24,6 @@ from Amadeus.symbolic_encoding import data_utils, decoding_utils
 from Amadeus.symbolic_encoding.data_utils import get_emb_total_size
 from Amadeus import model_zoo, trainer_accelerate as trainer
 from Amadeus.train_utils import NLLLoss4REMI, NLLLoss4CompoundToken, CosineAnnealingWarmUpRestarts, EncodecFlattenLoss, EncodecMultiClassLoss, CosineLRScheduler, adjust_prediction_order, DiffusionLoss4CompoundToken
-from Amadeus.encodec.data_utils import EncodecDataset
 from data_representation import vocab_utils
 from run_evaluation import main as run_evaluation