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