86cff17e7ef1b355e1ab4862ccd664ed4c856227,python/ray/util/sgd/torch/torch_trainer.py,TorchTrainer,get_model,#TorchTrainer#,498

Before Change

        if len(state["models"]) == 1:
            models.load_state_dict(state["models"][0])
        else:
            for model, state_dict in zip(models, state["models"]):
                model.load_state_dict(state_dict)
        return models

    def state_dict(self):
        return self.local_worker.get_state()

After Change

        unwrapped = []
        for model in self.local_worker.models:
            unwrapped += [model.module if hasattr(model, "module") else model]
        if len(unwrapped) == 1:
            return unwrapped[0]
        return unwrapped

Instances