diff --git a/demo/guide-python/feature_weights.py b/demo/guide-python/feature_weights.py
index 006c2f608296..07a8719422c6 100644
--- a/demo/guide-python/feature_weights.py
+++ b/demo/guide-python/feature_weights.py
@@ -12,8 +12,8 @@
 def main(args):
     rng = np.random.RandomState(1994)
 
-    kRows = 512
-    kCols = 64
+    kRows = 1000
+    kCols = 10
 
     X = rng.randn(kRows, kCols)
     y = rng.randn(kRows)
@@ -31,6 +31,7 @@ def main(args):
     featue_map = bst.get_fscore()
     # feature zero has 0 weight
     assert featue_map.get('f0', None) is None
+    assert max(featue_map.values()) == featue_map.get('f9')
 
     if args.plot:
         xgboost.plot_importance(bst)