admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-08-08 17:39:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix predictor attention dropout and drop its dead parameters

Browse Source
This commit is contained in:
Maxime Ellerbach
2026-08-02 23:05:36 +00:00
parent 80027e4460
commit 54eda5048b
5 changed files with 43 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/policies/vla_jepa/test_action_head.py
+2 -3
View File
@@ -53,13 +53,12 @@ def test_action_head_sample_time_range(action_dim: int, state_dim: int, action_h
def test_action_head_build_inputs_shape(action_dim: int, state_dim: int, action_horizon: int) -> None:
config = make_config(action_dim=action_dim, state_dim=state_dim, action_horizon=action_horizon)
head = VLAJEPAActionHead(config, cross_attention_dim=QWEN_HIDDEN_SIZE)
conditioning = torch.randn(2, 4, QWEN_HIDDEN_SIZE)
actions = torch.randn(2, action_horizon, action_dim)
timesteps = torch.randint(0, 100, (2,))
state = torch.randn(2, state_dim) if state_dim > 0 else None
out_with = head._build_inputs(conditioning, actions, state, timesteps)
out_none = head._build_inputs(conditioning, actions, None, timesteps)
out_with = head._build_inputs(actions, state, timesteps)
out_none = head._build_inputs(actions, None, timesteps)
assert out_with.ndim == 3 and out_none.ndim == 3
if state_dim > 0: