Update graph_decoder/diffusion_model.py

graph_decoder/diffusion_model.py

@@ -221,7 +221,6 @@ class GraphDiT(nn.Module):
         return noisy_data
 
     @torch.no_grad()
-    # @spaces.GPU(duration=60)
     def generate(
         self,
         properties,
@@ -313,7 +312,8 @@ class GraphDiT(nn.Module):
 
     def check_valid(self, smiles):
         return check_valid(smiles)
-    
+        
+    @spaces.GPU(duration=60)
     def sample_p_zs_given_zt(
         self, s, t, X_t, E_t, properties, node_mask, guide_scale, device
     ):