fix tdm remain layer

alibaba · Oct 8, 2024 · f1e8827 · f1e8827
1 parent 439d8a7
commit f1e8827
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/tzrec/datasets/dataset.py b/tzrec/datasets/dataset.py
@@ -128,7 +128,7 @@ def _expand_tdm_sample(
         )
         remain_layer.sort()
     else:
-        remain_layer = list(range(tree_level - 1))
+        remain_layer = np.array(range(tree_level - 1))
 
     num_remain_layer_neg = (
         sum([layer_num_sample[i] for i in remain_layer]) + layer_num_sample[-1]

diff --git a/tzrec/datasets/dataset_test.py b/tzrec/datasets/dataset_test.py
@@ -499,7 +499,7 @@ def _childern(code):
                 fg_encoded=True,
                 label_fields=["label"],
                 negative_sampler=sampler_pb2.TDMSampler(
-                    input_input_path=node.name,
+                    item_input_path=node.name,
                     edge_input_path=edge.name,
                     predict_edge_input_path=predict_edge.name,
                     attr_fields=["tree_level", "int_a", "float_b", "str_c"],