diff --git a/python/examples/criteo_ctr_with_cube/test_client.py b/python/examples/criteo_ctr_with_cube/test_client.py
index de205ebc68af02e8dd978da51a4c43bef0cec0d4..ca752b763e067b6a73e28c1d2ab9f58b9b98ba5d 100755
--- a/python/examples/criteo_ctr_with_cube/test_client.py
+++ b/python/examples/criteo_ctr_with_cube/test_client.py
@@ -40,7 +40,7 @@ for ei in range(10000):
     for i in range(1, 27):
         feed_dict["embedding_{}.tmp_0".format(i - 1)] = data[0][i]
     fetch_map = client.predict(feed=feed_dict, fetch=["prob"])
-    prob_list.append(fetch_map['prob'][1])
+    prob_list.append(fetch_map['prob'][0][1])
     label_list.append(data[0][-1][0])
 
 print(auc(label_list, prob_list))