# \[Tune\] There is a need to add resetup method in trainable class?

**URL:** <https://discuss.ray.io/t/tune-there-is-a-need-to-add-resetup-method-in-trainable-class/3119>\
**Category:** Ray Tune\
**Created:** [August 3, 2021, 1:35am UTC](https://discuss.ray.io/t/tune-there-is-a-need-to-add-resetup-method-in-trainable-class/3119 "2021-08-03T01:35:52Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![JUstForFUN](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/justforfun/32/1298_2.png) [@JUstForFUN](https://discuss.ray.io/u/JUstForFUN)\
**Post date:** [August 3, 2021, 1:35am UTC](https://discuss.ray.io/t/tune-there-is-a-need-to-add-resetup-method-in-trainable-class/3119/1 "2021-08-03T01:35:52Z")

</div>

Hi! I’m using the tune to find the hyper-parameters in my model. I just want to train the model with the same initialized parameters, and I want to reuse the actors to speed up my training process. I also use the **wandb** to visualize the every trial. I just implement my code as follows:

```
from ray.tune import Trainable
from ray.tune.integration.wandb import WandbTrainableMixin
from copy import deepcopy
import wandb

class My_Train(WandbTrainableMixin, Trainable):
    def setup(self, config):
        # I hope to restart the model on the same model parameters
        # It's time consuming to re-initialize the model
        # So I save the model state dict in self.model_state_dict

        self.model = ...
        self.model_state_dict = deepcopy(self.model.state_dict())

        self.hyper_params = config['hyper_params']
        self.epoch = 0

    def step(self):
        avg_loss = 0.0
        for batch_idx, batch in enumerate(self.dataloader):
            loss = self.model(batch)
            avg_loss = (avg_loss * batch_idx + loss) / (batch_idx + 1)

        wandb.log({'loss': avg_loss, 'epoch':self.epoch})

    def reset_config(self, new_config):
        # In order to restart the same initialized model with a different config
        self.hyper_params = new_config['hyper_params']
        self.model.load_state_dict(self.model_state_dict)
        self.epoch = 0

```

But I found that the latter trial’s {‘loss’: avg\_loss, ‘epoch’:self.epoch} will overwrite the former trial’s rather than create a new wandb init.

I know I can solve this problem by not using the reuse\_actors method. But I really want to know whether there is a method to solve this problem, thanks a lot!!

---

<div class="post-metadata">

**Author:** ![kai](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/kai/32/3380_2.png) [@kai](https://discuss.ray.io/u/kai)\
**Post date:** [August 3, 2021, 10:14am UTC](https://discuss.ray.io/t/tune-there-is-a-need-to-add-resetup-method-in-trainable-class/3119/2 "2021-08-03T10:14:05Z")

</div>

Are you using Wandb just to log the trial results? If so, you might be able to just use a `WandbLoggerCallback` instead of the mixin.

See [External library integrations (tune.integration) — Ray v2.0.0.dev0](https://docs.ray.io/en/master/tune/api_docs/integration.html#weights-and-biases-tune-integration-wandb)

The mixin calls `wandb.init` when the trainable is constructed - however, you should be able to re-initialize wandb in your `reset_config` method.

```auto
self.wandb = self._wandb.init(**wandb_init_kwargs)

```

see the source code here: [ray.tune.integration.wandb — Ray v2.0.0.dev0](https://docs.ray.io/en/master/_modules/ray/tune/integration/wandb.html)