diff --git a/codebook_features/train_codebook.py b/codebook_features/train_codebook.py
index 451e425..fcb8902 100755
--- a/codebook_features/train_codebook.py
+++ b/codebook_features/train_codebook.py
@@ -11,6 +11,7 @@
 import torch
 import transformers
 import wandb
+import pathlib
 
 from codebook_features import models, run_clm
 from codebook_features import trainer as cb_trainer