add safetensors

eval.py

@@ -6,6 +6,7 @@ from stable_baselines3.common.evaluation import evaluate_policy
 from stable_baselines3.common.atari_wrappers import AtariWrapper
 from stable_baselines3.common.env_util import make_vec_env
 from stable_baselines3.common.vec_env import VecFrameStack
+from safetensors.torch import save_model
 
 if __name__ == "__main__":
 
@@ -18,10 +19,10 @@ if __name__ == "__main__":
     
     agent = PPO.load("model.zip")
 
-    print(agent.policy)
-    total_params = sum(p.numel() for p in agent.policy.parameters())
-    print("\nTotal number of parameters: {:,}".format(total_params))
-
+    # Save the state dict as a safetensors file
+    policy_state_dict = agent.policy.state_dict()
+    save_model(agent.policy, "model.safetensors")
+    
     mean_reward, std_reward = evaluate_policy(agent, env, n_eval_episodes=n_eval_episodes, deterministic=deterministic)
     
     print(f"reward : {mean_reward} +/- {std_reward}")

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:32f14418431437b7d5f5a4e2ae80b4454939ccaab5cf811387a30409082e8baa
+size 8978048

requirements.txt

@@ -3,4 +3,5 @@ ale_py==0.10.1
 gymnasium[atari]
 torch==2.4.0
 stable_baselines3==2.4.1
-opencv-python==25.0.1
+opencv-python==25.0.1
+safetensors==0.5.3