diff --git a/example/train_small_criteo_demo.py b/example/train_small_criteo_demo.py
index e28106a..de619f2 100644
--- a/example/train_small_criteo_demo.py
+++ b/example/train_small_criteo_demo.py
@@ -24,7 +24,7 @@
     # TODO: Hyper Parameters
     file = 'data/criteo/train.txt'
     read_part = True
-    sample_num = 5000000
+    sample_num = 50000
     test_size = 0.2
 
     model_params = {
diff --git a/reclearn/models/ranking/deepfm.py b/reclearn/models/ranking/deepfm.py
index dea85eb..2e91d92 100644
--- a/reclearn/models/ranking/deepfm.py
+++ b/reclearn/models/ranking/deepfm.py
@@ -54,10 +54,10 @@ def call(self, inputs):
 		sparse_inputs = index_mapping(inputs, self.map_dict)
 		wide_inputs = {'sparse_inputs': sparse_inputs,
 					   'embed_inputs': tf.reshape(sparse_embed, shape=(-1, self.field_num, self.embed_dim))}
-		wide_outputs = self.fm(wide_inputs)  # (batch_size, 1)
+		wide_outputs = tf.reshape(self.fm(wide_inputs), [-1, 1])  # (batch_size, 1)
 		# deep
 		deep_outputs = self.mlp(sparse_embed)
-		deep_outputs = self.dense(deep_outputs)  # (batch_size, 1)
+		deep_outputs = tf.reshape(self.dense(deep_outputs), [-1, 1])  # (batch_size, 1)
 		# outputs
 		outputs = tf.nn.sigmoid(tf.add(wide_outputs, deep_outputs))
 		return outputs