Update graph_decoder/diffusion_model.py

graph_decoder/diffusion_model.py

@@ -1,5 +1,3 @@
-import spaces
-
 import os 
 import yaml
 import json
@@ -57,10 +55,6 @@ class GraphDiT(nn.Module):
             ydim=self.ydim,
         )
         self.model_dtype = model_dtype
-        # self.device = next(self.denoiser.parameters()).device
-
-        # model_params = torch.load(model_params_path, map_location='cpu')
-        # self.denoiser.load_state_dict(model_params)
         
         self.noise_schedule = utils.PredefinedNoiseScheduleDiscrete(
             dm_cfg.diffusion_noise_schedule, timesteps=dm_cfg.diffusion_steps
@@ -221,7 +215,7 @@ class GraphDiT(nn.Module):
         }
         return noisy_data
 
-    @torch.no_grad()
+    # @torch.no_grad()
     def generate(
         self,
         properties,