diff --git a/demo/image_classification/api_v2_train.py b/demo/image_classification/api_v2_train.py
index 44a8db39414f3b924c7fe837b6e250053a7bd3f9..e6e430724205754577fcc2bc1857b59fb494efec 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):