86cff17e7ef1b355e1ab4862ccd664ed4c856227,python/ray/util/sgd/torch/torch_runner.py,TorchRunner,_get_model_state_dicts,#TorchRunner#,227

Before Change

        // resume training while saving intermediate checkpoints.
        cpu_state_dicts = []
        for model in self.models:
            state_dict = model.state_dict()
            cpu_state_dicts += [{k: v.cpu() for k, v in state_dict.items()}]
        return cpu_state_dicts

    def _set_model_state_dicts(self, models_state_dicts):

After Change

        self.training_operator._set_timers(self.timers)

    def _get_model_state_dicts(self):
        return [model.state_dict() for model in self.models]

    def _set_model_state_dicts(self, models_state_dicts):
        for model, state_dict in zip(self.models, models_state_dicts):
            model.load_state_dict(state_dict)

Instances