From 494aa576b25cc092bec6970fcd087e41b2dcc1aa Mon Sep 17 00:00:00 2001 From: Pepijn Date: Thu, 11 Sep 2025 09:18:20 +0200 Subject: [PATCH] fix push to hub test --- test_pi0_pi05_hub.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test_pi0_pi05_hub.py b/test_pi0_pi05_hub.py index 522b32b2e..d737bc561 100644 --- a/test_pi0_pi05_hub.py +++ b/test_pi0_pi05_hub.py @@ -84,7 +84,8 @@ def test_hub_loading(model_id="pepijn223/pi0_base_fp32", model_name="PI0"): print(f" - State dimension: {policy.config.state_dim}") print(f" - Action horizon: {policy.config.action_horizon}") print(f" - Tokenizer max length: {policy.config.tokenizer_max_length}") - print(f" - discrete_state_input: {policy.config.discrete_state_input}") + if model_name == "PI0.5": + print(f" - discrete_state_input: {policy.config.discrete_state_input}") print(f" - Device: {device}") print(f" - Dtype: {next(policy.parameters()).dtype}")