From c98dba28eb1259c532cf787b26486835156882a9 Mon Sep 17 00:00:00 2001
From: raushan <raushan@huggingface.co>
Date: Tue, 14 Apr 2026 15:22:09 +0200
Subject: [PATCH] seems like a weird numerical precision issue, lets check in
 runners

---
 src/lerobot/policies/pi0_fast/modeling_pi0_fast.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/lerobot/policies/pi0_fast/modeling_pi0_fast.py b/src/lerobot/policies/pi0_fast/modeling_pi0_fast.py
index cd82838ea..d049fa00d 100644
--- a/src/lerobot/policies/pi0_fast/modeling_pi0_fast.py
+++ b/src/lerobot/policies/pi0_fast/modeling_pi0_fast.py
@@ -260,6 +260,8 @@ class PI0FastPaliGemma(nn.Module):
             image = image.to(torch.float32)
         image_outputs = self.paligemma.model.get_image_features(image)
         features = image_outputs.pooler_output
+        norm = 2048**0.5
+        features = features / norm * norm
         if features.dtype != out_dtype:
             features = features.to(out_dtype)
         return features