File size: 977 Bytes
96767e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 | diff --git a/src/lerobot/policies/pi052/modeling_pi052.py b/src/lerobot/policies/pi052/modeling_pi052.py
index 1a57dc4..b0ec5c5 100644
--- a/src/lerobot/policies/pi052/modeling_pi052.py
+++ b/src/lerobot/policies/pi052/modeling_pi052.py
@@ -176,8 +176,10 @@ class PI05Pytorch(PI05PytorchBase): # see openpi `PI0Pytorch`
finally:
self._lang_causal_marks = None
- def embed_prefix(self, images, img_masks, tokens, masks):
- prefix_embs, prefix_pad, prefix_att = super().embed_prefix(images, img_masks, tokens, masks)
+ def embed_prefix(self, images, img_masks, tokens, masks, states=None, state_masks=None):
+ prefix_embs, prefix_pad, prefix_att = super().embed_prefix(
+ images, img_masks, tokens, masks, states, state_masks
+ )
marks = getattr(self, "_lang_causal_marks", None)
if marks is not None:
prefix_att = _apply_causal_language_marks(prefix_att, marks.to(prefix_att.device))
|