b708d24f8a93d823d727435c24cbcd5713252640,catalyst/rl/core/trainer.py,TrainerSpec,_update_sampler_weights,#TrainerSpec#,139

Before Change

    def _update_sampler_weights(self):
        if self.epoch % self.weights_sync_period == 0:
            state_dict = self.algorithm.__dict__[
                self._weights_sync_mode].state_dict()
            state_dict = {
                k: v.detach().cpu().numpy()
                for k, v in state_dict.items()
            }

After Change

    def _update_sampler_weights(self):
        if self.epoch % self.weights_sync_period == 0:
            checkpoint = self.algorithm.pack_checkpoint(with_optimizer=False)
            for key in checkpoint:
                checkpoint[key] = {
                    k: v.detach().cpu().numpy()
                    for k, v in checkpoint[key].items()
                }

            self.db_server.save_checkpoint(
                checkpoint=checkpoint,
                epoch=self.epoch
            )

Instances