diff --git a/demo/mnist/api_train.py b/demo/mnist/api_train.py
index ae00e9f54c3a0e7da65c7738931ce3345bca98b2..c6993eda38bd022de0a923e3d530fd41535830d3 100644
--- a/demo/mnist/api_train.py
+++ b/demo/mnist/api_train.py
@@ -99,7 +99,7 @@ def main():
     # DataProvider Converter is a utility convert Python Object to Paddle C++
     # Input. The input format is as same as Paddle's DataProvider.
     converter = DataProviderConverter(
-        input_types=[dp.dense_vector(784), dp.integer_value(10)])
+        input_types=[images.data_type, label.data_type])
 
     train_file = './data/raw_data/train'
     test_file = './data/raw_data/t10k'
diff --git a/demo/mnist/api_train_v2.py b/demo/mnist/api_train_v2.py
index b46cf1c8700ba8de2f9e5f8af3320ed6e1d9f7a9..6bbc3a4ab35a64b42c60eb14beca4ff2b6e471b7 100644
--- a/demo/mnist/api_train_v2.py
+++ b/demo/mnist/api_train_v2.py
@@ -51,8 +51,8 @@ def main():
                   batch_size=32,  # batch size should be refactor in Data reader
                   data_types={  # data_types will be removed, It should be in
                       # network topology
-                      'pixel': dense_vector(784),
-                      'label': integer_value(10)
+                      'pixel': images.data_type,
+                      'label': label.data_type
                   })
 
 
diff --git a/python/paddle/v2/layer.py b/python/paddle/v2/layer.py
index 511b3e7457c261c2aed69c3a55ab34cd06bd4a78..f2bc47946b6df877658f409bb7583f02ec62493b 100644
--- a/python/paddle/v2/layer.py
+++ b/python/paddle/v2/layer.py
@@ -168,15 +168,15 @@ class DataLayerV2(Layer):
     def __init__(self, name, data_type, **kwargs):
         assert isinstance(data_type, dp.InputType)
 
+        self.data_type = data_type
         self.__method_name__ = 'data_layer'
         self.__kwargs__ = kwargs
-        self.__data_size__ = data_type.dim
 
         super(DataLayerV2, self).__init__(name=name, parent_layers=dict())
 
     def to_proto_impl(self, **kwargs):
         args = dict()
-        args['size'] = self.__data_size__
+        args['size'] = self.data_type.dim
         for each in kwargs:
             args[each] = kwargs[each]
         for each in self.__kwargs__: