diff --git a/train.py b/train.py
index 3bd4c9f..3c94270 100644
--- a/train.py
+++ b/train.py
@@ -3,6 +3,7 @@ import tensorflow as tf
 from PIL import Image
 import tqdm
 import scipy.ndimage
+import os
 
 from model import Model
 import load