# How to save model during tuning

**URL:** <https://discuss.ray.io/t/how-to-save-model-during-tuning/13337>\
**Category:** Checkpointing, Restoring\
**Created:** [January 8, 2024, 3:42am UTC](https://discuss.ray.io/t/how-to-save-model-during-tuning/13337 "2024-01-08T03:42:06Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![Keik1999](https://avatars.discourse-cdn.com/v4/letter/k/c6cbf5/32.png) [@Keik1999](https://discuss.ray.io/u/Keik1999)\
**Post date:** [January 8, 2024, 3:42am UTC](https://discuss.ray.io/t/how-to-save-model-during-tuning/13337/1 "2024-01-08T03:42:07Z")

</div>

i want to save the model during tune process when i get a new best reward, not only save the model after finish the whole tune.  
but i do not know how to write the code.  
i tried to use the callback function but i do not the the save function which can be used  
here is my code

```auto
import ray
from ray import train, tune
from ray.rllib.algorithms import Algorithm
from ray.tune import Callback
from ray.tune.registry import get_trainable_cls

class SaveModelOnNewMaxRewardCallback(Callback):
    def __init__ (self, max_reward_key="episode_reward_max", model_save_path="model_checkpoint"):
        self.max_reward_key = max_reward_key
        self.model_save_path = model_save_path
        self.best_reward = float("-inf")

    def on_trial_result(self, iteration, trials, trial, result, **info):
        new_max_reward = result.get(self.max_reward_key, None)
        if new_max_reward is not None and new_max_reward > self.best_reward:
            self.best_reward = new_max_reward
            # save the model
            trial.save_checkpoint(self.model_save_path)

ray.init()

config = (
    get_trainable_cls('PPO')
    .get_default_config()
    .environment("CartPole-v1")
    .framework("torch")
    .rollouts(num_rollout_workers=5)
)
# 创建回调函数实例
save_model_callback = SaveModelOnNewMaxRewardCallback()
# ``Tuner.fit()`` allows setting a custom log directory (other than ``~/ray-results``)
tuner = ray.tune.Tuner(
    "PPO",
    param_space=config,
    run_config=train.RunConfig(
        stop={"episode_reward_mean": 50},
        checkpoint_config=train.CheckpointConfig(checkpoint_at_end=True),
        callbacks=[save_model_callback],
    )
)

results = tuner.fit()

# Get the best result based on a particular metric.
best_result = results.get_best_result(metric="episode_reward_mean", mode="max")

# Get the best checkpoint corresponding to the best result.
best_checkpoint = best_result.checkpoint

algo = Algorithm.from_checkpoint("model_checkpoint")

```

there is no trial.save\_checkpoint(self.model\_save\_path) function  
thanks a lot