diff --git a/diffusion_schedulers/scheduling_flow_matching.py b/diffusion_schedulers/scheduling_flow_matching.py
index 3345304..b8c1507 100644
--- a/diffusion_schedulers/scheduling_flow_matching.py
+++ b/diffusion_schedulers/scheduling_flow_matching.py
@@ -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,
             )