| 58 | |
| 59 | |
| 60 | def load_checkpoint(args, model, optimizer): |
| 61 | CHECKPOINT_PATH = os.path.join(args.local_workspace, 'checkpoint_latest.pth') |
| 62 | checkpoint = torch.load(CHECKPOINT_PATH, map_location=torch.device('cpu')) |
| 63 | model.load_state_dict(checkpoint['model']) |
| 64 | optimizer.load_state_dict(checkpoint['optimizer']) |
| 65 | start_epoch = checkpoint['epoch'] |
| 66 | print(colored('[MODEL]: ', 'yellow') + 'Loading the model at epoch: {:d}'.format(start_epoch)) |
| 67 | |
| 68 | return start_epoch + 1 |
| 69 | |
| 70 | def train(rank, world_size, args): |
| 71 | print(f"Training on rank {rank}.") |