| 543 | state_dict.update(shard_state) |
| 544 | |
| 545 | class _StateDictModel: |
| 546 | def __init__(self, config, state_dict): |
| 547 | self.config = config |
| 548 | self._state_dict = state_dict |
| 549 | |
| 550 | def state_dict(self): |
| 551 | return self._state_dict |
| 552 | |
| 553 | model = _StateDictModel(config_obj, state_dict) |
| 554 |