(self, checkpoint_path, on_gpu)
| 659 | return did_restore |
| 660 | |
| 661 | def restore(self, checkpoint_path, on_gpu): |
| 662 | checkpoint = torch.load(checkpoint_path, map_location='cpu') |
| 663 | |
| 664 | # load model state |
| 665 | model = self.get_model() |
| 666 | |
| 667 | # load the state_dict on the model automatically |
| 668 | model.load_state_dict(checkpoint['state_dict'], strict=False) |
| 669 | if on_gpu: |
| 670 | model.cuda(self.root_gpu) |
| 671 | # load training state (affects trainer only) |
| 672 | self.restore_training_state(checkpoint) |
| 673 | model.global_step = self.global_step |
| 674 | del checkpoint |
| 675 | |
| 676 | try: |
| 677 | if dist.is_initialized() and dist.get_rank() > 0: |
| 678 | return |
| 679 | except Exception as e: |
| 680 | print(e) |
| 681 | return |
| 682 | |
| 683 | def restore_training_state(self, checkpoint): |
| 684 | """ |
no test coverage detected