| | |
|---|
| | | if use_ddp: |
|---|
| | | model = model.cuda(local_rank) |
|---|
| | | model = DDP(model, device_ids=[local_rank], |
|---|
| | | find_unused_parameters=kwargs.get("train_conf", {}).get("find_unused_parameters", False)) |
|---|
| | | find_unused_parameters=kwargs.get("train_conf", {}).get("find_unused_parameters", True)) |
|---|
| | | elif use_fsdp: |
|---|
| | | # model = FSDP(model).cuda(local_rank) |
|---|
| | | |
|---|
| | |
|---|
| | | |
|---|
| | | |
|---|
| | | if trainer.rank == 0: |
|---|
| | | average_checkpoints(trainer.output_dir, trainer.avg_nbest_model, trainer.val_acc_list) |
|---|
| | | average_checkpoints(trainer.output_dir, trainer.avg_nbest_model) |
|---|
| | | |
|---|
| | | trainer.close() |
|---|
| | | |
|---|
