fix(g05): autocast training forward

src/lerobot/policies/g05/modeling_g05.py

@@ -389,6 +389,12 @@ class G05Policy(PreTrainedPolicy):
 
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict[str, Any] | None]:
         prepared = self._prepare_author_batch(batch)
+        device = next(self.backend.parameters()).device
+        with torch.autocast(
+            device_type=device.type,
+            dtype=torch.bfloat16,
+            enabled=self.config.model_weights_to_bf16 and device.type == "cuda",
+        ):
             result = self.backend(prepared)
         if isinstance(result, tuple) and len(result) == 2:
             loss, loss_dict = result

tests/policies/g05/test_g05.py

@@ -576,6 +576,28 @@ def test_forward_backward_update_and_save_reload(tmp_path: Path):
     torch.testing.assert_close(actual, expected)
 
 
+def test_training_forward_uses_policy_autocast_context(monkeypatch):
+    policy = G05Policy(_config(), backend=TinyG05Backend())
+    autocast_calls = []
+
+    class AutocastContext:
+        def __enter__(self):
+            return None
+
+        def __exit__(self, exc_type, exc_value, traceback):
+            return False
+
+    def track_autocast(**kwargs):
+        autocast_calls.append(kwargs)
+        return AutocastContext()
+
+    monkeypatch.setattr(torch, "autocast", track_autocast)
+
+    policy(_policy_batch("train"))
+
+    assert autocast_calls == [{"device_type": "cpu", "dtype": torch.bfloat16, "enabled": False}]
+
+
 def test_save_pretrained_copies_required_gated_sidecars_portably(tmp_path: Path):
     source = tmp_path / "checkpoint"
     processor = source / "hf_processor"