| 633 | raw_state_dict = _load_raw_hf_state_dict(model_id, cast_to_bf16=cast_to_bf16) |
| 634 | |
| 635 | class _RawModelWrapper: |
| 636 | def __init__(self, state_dict, config): |
| 637 | self._state_dict = state_dict |
| 638 | self.config = config |
| 639 | |
| 640 | def state_dict(self): |
| 641 | return self._state_dict |
| 642 | |
| 643 | model = _RawModelWrapper(raw_state_dict, config_json) |
| 644 | else: |