funasr/bin/train_ds.py
@@ -84,9 +84,6 @@ dist.init_process_group( backend=kwargs.get("backend", "nccl"), init_method="env://", find_unused_parameters=kwargs.get("train_conf", {}).get( "find_unused_parameters", False ), ) torch.cuda.set_device(local_rank) @@ -137,7 +134,7 @@ **kwargs.get("train_conf"), ) model = trainer.warp_model(model) model = trainer.warp_model(model, **kwargs) kwargs["device"] = int(os.environ.get("LOCAL_RANK", 0)) trainer.device = int(os.environ.get("LOCAL_RANK", 0))
