diff --git a/train.py b/train.py
index b5085c2..25e7743 100644
--- a/train.py
+++ b/train.py
@@ -12,6 +12,7 @@
 os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
 
 import jax
+import json
 import numpy as np
 from pipe.pipe import SpectrogramPipe
 from pyspark.ml.feature import StringIndexer, IndexToString