| 327 | |
| 328 | |
| 329 | def auto_resume_helper(output_dir): |
| 330 | checkpoints = os.listdir(output_dir) |
| 331 | checkpoints = [ckpt for ckpt in checkpoints if ckpt.endswith('pth')] |
| 332 | print(f'All checkpoints founded in {output_dir}: {checkpoints}') |
| 333 | if len(checkpoints) > 0: |
| 334 | latest_checkpoint = max( |
| 335 | [os.path.join(output_dir, d) for d in checkpoints], |
| 336 | key=os.path.getmtime) |
| 337 | print(f'The latest checkpoint founded: {latest_checkpoint}') |
| 338 | resume_file = latest_checkpoint |
| 339 | else: |
| 340 | resume_file = None |
| 341 | return resume_file |
| 342 | |
| 343 | |
| 344 | def reduce_tensor(tensor): |