From 3b87bd5609d45e2eb6c5523eeaef6b47e1a41436 Mon Sep 17 00:00:00 2001 From: Mufei Li Date: Sat, 30 Apr 2022 14:52:11 +0800 Subject: [PATCH] Update jtvae.py (#178) --- python/dgllife/model/model_zoo/jtvae.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/dgllife/model/model_zoo/jtvae.py b/python/dgllife/model/model_zoo/jtvae.py index 563e8340..af386dcf 100644 --- 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