diff --git a/demo/image_classification/api_v2_train.py b/demo/image_classification/api_v2_train.py
index 44a8db3941..e6e4307242 100644
--- a/demo/image_classification/api_v2_train.py
+++ b/demo/image_classification/api_v2_train.py
@@ -14,6 +14,7 @@
 
 from api_v2_vgg import resnet_cifar10
 from api_v2_resnet import vgg_bn_drop
+import paddle.v2 as paddle
 
 
 def event_handler(event):