Copy a model state_dict to cpu. Args: state_dict (OrderedDict): Model weights on GPU. Returns: OrderedDict: Model weights on GPU.
(state_dict)
| 357 | |
| 358 | |
| 359 | def weights_to_cpu(state_dict): |
| 360 | """Copy a model state_dict to cpu. |
| 361 | |
| 362 | Args: |
| 363 | state_dict (OrderedDict): Model weights on GPU. |
| 364 | |
| 365 | Returns: |
| 366 | OrderedDict: Model weights on GPU. |
| 367 | """ |
| 368 | state_dict_cpu = OrderedDict() |
| 369 | for key, val in state_dict.items(): |
| 370 | state_dict_cpu[key] = val.cpu() |
| 371 | return state_dict_cpu |
| 372 | |
| 373 | |
| 374 | def _save_to_state_dict(module, destination, prefix, keep_vars): |
no outgoing calls
no test coverage detected