mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-29 12:39:41 +00:00
refactor(rtc): share Pi05 training helpers
This commit is contained in:
@@ -56,8 +56,8 @@ def test_create_sinusoidal_pos_embedding_matches_openpi_formula():
|
||||
def test_create_sinusoidal_pos_embedding_validation():
|
||||
with pytest.raises(ValueError, match="divisible by 2"):
|
||||
create_sinusoidal_pos_embedding(torch.zeros(2), 7, 4e-3, 4.0, device=torch.device("cpu"))
|
||||
with pytest.raises(ValueError, match="batch_size"):
|
||||
create_sinusoidal_pos_embedding(torch.zeros(2, 2), 8, 4e-3, 4.0, device=torch.device("cpu"))
|
||||
with pytest.raises(ValueError, match="must have shape"):
|
||||
create_sinusoidal_pos_embedding(torch.zeros(2, 2, 2), 8, 4e-3, 4.0, device=torch.device("cpu"))
|
||||
|
||||
|
||||
def test_make_att_2d_masks_docstring_cases():
|
||||
|
||||
@@ -16,13 +16,13 @@ from torch import nn
|
||||
pytest.importorskip("transformers")
|
||||
|
||||
from lerobot.policies.pi05.modeling_pi05 import ( # noqa: E402
|
||||
_build_flow_matching_inputs,
|
||||
_prepare_trained_rtc_prefix,
|
||||
create_sinusoidal_pos_embedding,
|
||||
)
|
||||
from lerobot.policies.pi052.configuration_pi052 import PI052Config # noqa: E402
|
||||
from lerobot.policies.pi052.modeling_pi052 import ( # noqa: E402
|
||||
PI05Pytorch as PI052Pytorch,
|
||||
_build_flow_matching_inputs,
|
||||
_flow_loss_per_sample,
|
||||
_reduce_flow_loss,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user