Update jtvae.py (#178)

awslabs · Apr 30, 2022 · 3b87bd5 · 3b87bd5
1 parent e98aee5
commit 3b87bd5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/dgllife/model/model_zoo/jtvae.py b/python/dgllife/model/model_zoo/jtvae.py
@@ -270,7 +270,7 @@ def forward(self, tree_graphs, tree_vec):
 
         # Traverse the tree and predict on children
         for eid, p in dfs_order(tree_graphs, root_ids.to(dtype=tree_graphs.idtype)):
-            eid = eid.to(device)
+            eid = eid.to(device=device, dtype=tree_graphs.idtype)
             p = p.to(device=device, dtype=tree_graphs.idtype)
 
             # Message passing excluding the target