Update monai/networks/schedulers/ddim.py

Signed-off-by: ytl0623 <[email protected]>

Author: ytl0623

monai/networks/schedulers/ddim.py

@@ -120,7 +120,11 @@ def set_timesteps(self, num_inference_steps: int, device: str | torch.device | N
         if self.steps_offset < 0 or self.steps_offset >= self.num_train_timesteps:
             raise ValueError(f"`steps_offset`: {self.steps_offset} must be in range [0, {self.num_train_timesteps}).")
 
-        self.timesteps = torch.linspace((self.num_train_timesteps - 1) - self.steps_offset, 0, num_inference_steps, device=device).round().long()
+        self.timesteps = (
+            torch.linspace((self.num_train_timesteps - 1) - self.steps_offset, 0, num_inference_steps, device=device)
+            .round()
+            .long()
+        )
         self.timesteps += self.steps_offset
 
     def _get_variance(self, timestep: int, prev_timestep: int) -> torch.Tensor:

monai/networks/schedulers/ddpm.py

@@ -122,8 +122,9 @@ def set_timesteps(self, num_inference_steps: int, device: str | torch.device | N
             )
 
         self.num_inference_steps = num_inference_steps
-        self.timesteps = torch.linspace(self.num_train_timesteps - 1, 0, self.num_inference_steps, device=device).round().long()
-        
+        self.timesteps = (
+            torch.linspace(self.num_train_timesteps - 1, 0, self.num_inference_steps, device=device).round().long()
+        )
 
     def _get_mean(self, timestep: int, x_0: torch.Tensor, x_t: torch.Tensor) -> torch.Tensor:
         """