Update qm9_mol_gen_ddpm.py

Author: amorehead

src/models/qm9_mol_gen_ddpm.py

@@ -648,6 +648,7 @@ def optimize(
         id_from: int = 0,
         chain_viz_batch_element_idx: int = 0,
         name: str = os.sep + "chain",
+        norm_with_original_timesteps: bool = False,
         verbose: bool = True
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         # context-conditioning
@@ -665,7 +666,8 @@ def optimize(
             context=context,
             device=self.device,
             num_timesteps=num_timesteps,
-            return_frames=return_frames
+            return_frames=return_frames,
+            norm_with_original_timesteps=norm_with_original_timesteps,
         )
 
         # visualize optimized samples