diff --git a/demo/image_classification/api_v2_train.py b/demo/image_classification/api_v2_train.py
index 7134fa61e861b9f41f49f44f55f3b19de96f5850..53cffa6fb4e8b2e19725f4f44bf7b9ffffb25232 100644
--- a/demo/image_classification/api_v2_train.py
+++ b/demo/image_classification/api_v2_train.py
@@ -16,7 +16,7 @@ import sys
 
 import paddle.v2 as paddle
 
-from api_v2_resnet import resnet_cifar10
+from api_v2_vgg import vgg_bn_drop
 
 
 def main():
@@ -31,9 +31,9 @@ def main():
 
     # Add neural network config
     # option 1. resnet
-    net = resnet_cifar10(image, depth=32)
+    # net = resnet_cifar10(image, depth=32)
     # option 2. vgg
-    # net = vgg_bn_drop(image)
+    net = vgg_bn_drop(image)
 
     out = paddle.layer.fc(input=net,
                           size=classdim,