diff --git a/tests/data/save_policy_to_safetensors/aloha_act/actions.safetensors b/tests/data/save_policy_to_safetensors/aloha_act/actions.safetensors
index 9658087b..59dd45ca 100644
Binary files a/tests/data/save_policy_to_safetensors/aloha_act/actions.safetensors and b/tests/data/save_policy_to_safetensors/aloha_act/actions.safetensors differ
diff --git a/tests/data/save_policy_to_safetensors/aloha_act/grad_stats.safetensors b/tests/data/save_policy_to_safetensors/aloha_act/grad_stats.safetensors
index 462f618a..29e009e5 100644
Binary files a/tests/data/save_policy_to_safetensors/aloha_act/grad_stats.safetensors and b/tests/data/save_policy_to_safetensors/aloha_act/grad_stats.safetensors differ
diff --git a/tests/data/save_policy_to_safetensors/aloha_act/output_dict.safetensors b/tests/data/save_policy_to_safetensors/aloha_act/output_dict.safetensors
index d2fd1693..c0ef62d0 100644
Binary files a/tests/data/save_policy_to_safetensors/aloha_act/output_dict.safetensors and b/tests/data/save_policy_to_safetensors/aloha_act/output_dict.safetensors differ
diff --git a/tests/data/save_policy_to_safetensors/aloha_act/param_stats.safetensors b/tests/data/save_policy_to_safetensors/aloha_act/param_stats.safetensors
index b03a032d..eef8f73f 100644
Binary files a/tests/data/save_policy_to_safetensors/aloha_act/param_stats.safetensors and b/tests/data/save_policy_to_safetensors/aloha_act/param_stats.safetensors differ
diff --git a/tests/data/save_policy_to_safetensors/pusht_diffusion/actions.safetensors b/tests/data/save_policy_to_safetensors/pusht_diffusion/actions.safetensors
index 85c8baae..6a748a64 100644
Binary files a/tests/data/save_policy_to_safetensors/pusht_diffusion/actions.safetensors and b/tests/data/save_policy_to_safetensors/pusht_diffusion/actions.safetensors differ
diff --git a/tests/data/save_policy_to_safetensors/pusht_diffusion/grad_stats.safetensors b/tests/data/save_policy_to_safetensors/pusht_diffusion/grad_stats.safetensors
index 2086112c..a551c0e8 100644
Binary files a/tests/data/save_policy_to_safetensors/pusht_diffusion/grad_stats.safetensors and b/tests/data/save_policy_to_safetensors/pusht_diffusion/grad_stats.safetensors differ
diff --git a/tests/data/save_policy_to_safetensors/pusht_diffusion/output_dict.safetensors b/tests/data/save_policy_to_safetensors/pusht_diffusion/output_dict.safetensors
index a50402bf..1e4f4e66 100644
Binary files a/tests/data/save_policy_to_safetensors/pusht_diffusion/output_dict.safetensors and b/tests/data/save_policy_to_safetensors/pusht_diffusion/output_dict.safetensors differ
diff --git a/tests/data/save_policy_to_safetensors/pusht_diffusion/param_stats.safetensors b/tests/data/save_policy_to_safetensors/pusht_diffusion/param_stats.safetensors
index 148d9d13..afcc1c3a 100644
Binary files a/tests/data/save_policy_to_safetensors/pusht_diffusion/param_stats.safetensors and b/tests/data/save_policy_to_safetensors/pusht_diffusion/param_stats.safetensors differ
diff --git a/tests/data/save_policy_to_safetensors/xarm_tdmpc/actions.safetensors b/tests/data/save_policy_to_safetensors/xarm_tdmpc/actions.safetensors
index d38cecb4..c0962440 100644
Binary files a/tests/data/save_policy_to_safetensors/xarm_tdmpc/actions.safetensors and b/tests/data/save_policy_to_safetensors/xarm_tdmpc/actions.safetensors differ
diff --git a/tests/data/save_policy_to_safetensors/xarm_tdmpc/grad_stats.safetensors b/tests/data/save_policy_to_safetensors/xarm_tdmpc/grad_stats.safetensors
index 89b2d7b8..c183097f 100644
Binary files a/tests/data/save_policy_to_safetensors/xarm_tdmpc/grad_stats.safetensors and b/tests/data/save_policy_to_safetensors/xarm_tdmpc/grad_stats.safetensors differ
diff --git a/tests/data/save_policy_to_safetensors/xarm_tdmpc/output_dict.safetensors b/tests/data/save_policy_to_safetensors/xarm_tdmpc/output_dict.safetensors
index cb6ebea6..2e8b53a3 100644
Binary files a/tests/data/save_policy_to_safetensors/xarm_tdmpc/output_dict.safetensors and b/tests/data/save_policy_to_safetensors/xarm_tdmpc/output_dict.safetensors differ
diff --git a/tests/data/save_policy_to_safetensors/xarm_tdmpc/param_stats.safetensors b/tests/data/save_policy_to_safetensors/xarm_tdmpc/param_stats.safetensors
index 031da1f6..a865fc29 100644
Binary files a/tests/data/save_policy_to_safetensors/xarm_tdmpc/param_stats.safetensors and b/tests/data/save_policy_to_safetensors/xarm_tdmpc/param_stats.safetensors differ