diff --git a/u2net_train.py b/u2net_train.py
index 43d9a12..8f19491 100644
--- a/u2net_train.py
+++ b/u2net_train.py
@@ -1,3 +1,4 @@
+import os
 import torch
 import torchvision
 from torch.autograd import Variable