Copy a model state_dict to cpu. Args: state_dict (OrderedDict): Model weights on GPU. Returns: OrderedDict: Model weights on GPU.
(state_dict)
| 28 | |
| 29 | |
| 30 | def weights_to_cpu(state_dict): |
| 31 | """Copy a model state_dict to cpu. |
| 32 | |
| 33 | Args: |
| 34 | state_dict (OrderedDict): Model weights on GPU. |
| 35 | |
| 36 | Returns: |
| 37 | OrderedDict: Model weights on GPU. |
| 38 | """ |
| 39 | state_dict_cpu = OrderedDict() |
| 40 | for key, val in state_dict.items(): |
| 41 | state_dict_cpu[key] = val.cpu() |
| 42 | # Keep metadata in state_dict |
| 43 | state_dict_cpu._metadata = getattr(state_dict, '_metadata', OrderedDict()) |
| 44 | return state_dict_cpu |
| 45 | |
| 46 | |
| 47 | def save_checkpoint(model: torch.nn.Module, |
no test coverage detected
searching dependent graphs…