diff --git a/wall_x/model/action_head.py b/wall_x/model/action_head.py
index 18dab36..cd5bb89 100755
--- a/wall_x/model/action_head.py
+++ b/wall_x/model/action_head.py
@@ -307,7 +307,7 @@ class ActionProcessor(nn.Module):
             torch.Tensor: Sampled timesteps of shape [batch_size]
         """
         sample = self.beta_dist.sample([batch_size]).to(device=device, dtype=dtype)
-        time = (1 - sample) / self.s
+        time = (1 - sample) * self.s
         return time
 
     def proprioception_proj(
diff --git a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py b/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py
index 9b45b51..3d31ef7 100644
--- a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py
+++ b/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py
@@ -2046,7 +2046,7 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
             times = torch.linspace(
                 0,
                 1,
-                num_inference_timesteps,
+                num_inference_timesteps + 1,
                 device=inputs_embeds.device,
                 dtype=inputs_embeds.dtype,
             )