Update

awslabs · Oct 10, 2023 · e7ac4a1 · e7ac4a1
1 parent b14f302
commit e7ac4a1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/graphstorm/inference/ep_infer.py b/python/graphstorm/inference/ep_infer.py
@@ -141,7 +141,7 @@ def infer(self, loader, save_embed_path, save_prediction_path=None,
                 target_ntypes.add(etype[0])
                 target_ntypes.add(etype[2])
             # Only init the nid_shuffler when there is a node_id_mapping_file.
-            nid_shuffler = NodeIDShuffler(g, node_id_mapping_file, target_ntypes) \
+            nid_shuffler = NodeIDShuffler(g, node_id_mapping_file, list(target_ntypes)) \
                 if node_id_mapping_file else None
             shuffled_preds = {}
             for etype, pred in preds.items():