diff --git a/tests/data/save_policy_to_safetensors/aloha_act/actions.safetensors b/tests/data/save_policy_to_safetensors/aloha_act/actions.safetensors
new file mode 100644
index 00000000..59dd45ca
Binary files /dev/null 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
new file mode 100644
index 00000000..173ca251
Binary files /dev/null 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
new file mode 100644
index 00000000..c0ef62d0
Binary files /dev/null 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
new file mode 100644
index 00000000..64e31f66
Binary files /dev/null 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
new file mode 100644
index 00000000..6273500c
Binary files /dev/null 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
new file mode 100644
index 00000000..e809d744
Binary files /dev/null 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
new file mode 100644
index 00000000..1c4c983d
Binary files /dev/null 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
new file mode 100644
index 00000000..d1044048
Binary files /dev/null and b/tests/data/save_policy_to_safetensors/pusht_diffusion/param_stats.safetensors differ