for merge

PaddleRec/tdm/tdm_demo/args.py

@@ -97,7 +97,7 @@ def parse_args():
 
     # for infer
     model_g.add_arg("child_nums", int, 2, "child node of ancestor node")
-    model_g.add_arg("topK", int, 2, "best recall result nums")
+    model_g.add_arg("topK", int, 1, "best recall result nums")
 
     model_g = ArgumentGroup(
         parser, "path", "files path of data & model.")

PaddleRec/tdm/tdm_demo/infer_network.py

@@ -104,7 +104,7 @@ class TdmInferNet(object):
         )
 
         for layer_idx in range(self.first_layer_idx, self.max_layers):
-            if layer_idx == 0:
+            if layer_idx == self.first_layer_idx:
                 current_layer_node_num = len(self.first_layer_node)
             else:
                 current_layer_node_num = current_layer_node.shape[1] * \