Skip to content

Commit

Permalink
handle both tensor and numpy input in init_sigmas_for_each_stage
Browse files Browse the repository at this point in the history
  • Loading branch information
kijai committed Oct 12, 2024
1 parent 63982da commit 2b80e37
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion diffusion_schedulers/scheduling_flow_matching.py
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,8 @@ def init_sigmas_for_each_stage(self):
timesteps = np.linspace(
timestep_max, timestep_min, training_steps + 1,
)
self.timesteps_per_stage[i_s] = torch.from_numpy(timesteps[:-1])
self.timesteps_per_stage[i_s] = timesteps[:-1] if isinstance(timesteps, torch.Tensor) else torch.from_numpy(timesteps[:-1])

stage_sigmas = np.linspace(
1, 0, training_steps + 1,
)
Expand Down

0 comments on commit 2b80e37

Please sign in to comment.