Update graph_decoder/diffusion_model.py

graph_decoder/diffusion_model.py

@@ -181,11 +181,11 @@ class GraphDiT(nn.Module):
     def generate(
         self,
         properties,
-        device,
         guide_scale=1.,
         num_nodes=None,
         number_chain_steps=50,
     ):
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         properties = [float('nan') if x is None else x for x in properties]
         properties = torch.tensor(properties, dtype=torch.float).reshape(1, -1).to(device)
         batch_size = properties.size(0)