(args: APNamespace)
| 535 | |
| 536 | |
| 537 | def main_worker(args: APNamespace): |
| 538 | torch.manual_seed(args.seed) |
| 539 | if torch.cuda.is_available() and not args.cpu: |
| 540 | device = 'cuda' |
| 541 | autocast = True |
| 542 | elif torch.backends.mps.is_available() and not args.cpu: |
| 543 | device = 'mps' |
| 544 | autocast = False |
| 545 | else: |
| 546 | device = 'cpu' |
| 547 | autocast = False |
| 548 | training_agent = TrainingAgent( |
| 549 | config_path=args.config_path, |
| 550 | device=device, |
| 551 | output_path=args.output_path, |
| 552 | data_path=args.data_path, |
| 553 | resume=args.resume, |
| 554 | save_freq=args.save_freq, |
| 555 | checkpoint_path=args.checkpoint_path, |
| 556 | dist=args.dist, |
| 557 | autocast=autocast, |
| 558 | pretrained=args.pretrained) |
| 559 | print(f"Info: Pytorch device is set to {training_agent.device}") |
| 560 | training_agent.train() |
| 561 | |
| 562 | |
| 563 | if __name__ == "__main__": |
no test coverage detected