diff --git a/u2net_train.py b/u2net_train.py
index 40395e9..24466d4 100644
--- a/u2net_train.py
+++ b/u2net_train.py
@@ -11,6 +11,7 @@ import torchvision.transforms as standard_transforms
 
 import numpy as np
 import glob
+import os
 
 from data_loader import Rescale
 from data_loader import RescaleT