handle both tensor and numpy input in init_sigmas_for_each_stage

from upstream: https://github.com/jy0205/Pyramid-Flow/commit/bf10e7340cd1e3467c7b88053bc0f413eb1a6073#diff-8e1c08bf8ab5d587ee253dccfb631b170202adf0664272785a9b377975ffb8a5R145
2024-10-12 10:45:01 +03:00
parent 63982da076
commit 2b80e37555
1 changed files with 2 additions and 1 deletions
@@ -142,7 +142,8 @@ class PyramidFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
            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,
            )