(saved_dict: Dict[str, Any])
| 142 | |
| 143 | # Function to extract the state_dict from a saved checkpoint |
| 144 | def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]: |
| 145 | state_dict = saved_dict |
| 146 | if "model" in saved_dict.keys(): |
| 147 | state_dict = state_dict["model"] |
| 148 | if "module" in saved_dict.keys(): |
| 149 | state_dict = state_dict["module"] |
| 150 | if "state_dict" in saved_dict.keys(): |
| 151 | state_dict = state_dict["state_dict"] |
| 152 | return state_dict |
| 153 | |
| 154 | |
| 155 | # Function to update the state_dict with new key assignments in-place |
no outgoing calls
no test coverage detected