support cpu training, use cpu training on mac
This commit is contained in:
@@ -118,16 +118,16 @@ def main(args):
|
||||
os.environ["MASTER_ADDR"]="localhost"
|
||||
trainer: Trainer = Trainer(
|
||||
max_epochs=config["train"]["epochs"],
|
||||
accelerator="gpu",
|
||||
accelerator="gpu" if torch.cuda.is_available() else "cpu",
|
||||
# val_check_interval=9999999999999999999999,###不要验证
|
||||
# check_val_every_n_epoch=None,
|
||||
limit_val_batches=0,
|
||||
devices=-1,
|
||||
devices=-1 if torch.cuda.is_available() else 1,
|
||||
benchmark=False,
|
||||
fast_dev_run=False,
|
||||
strategy = "auto" if torch.backends.mps.is_available() else DDPStrategy(
|
||||
strategy = DDPStrategy(
|
||||
process_group_backend="nccl" if platform.system() != "Windows" else "gloo"
|
||||
), # mps 不支持多节点训练
|
||||
) if torch.cuda.is_available() else "auto",
|
||||
precision=config["train"]["precision"],
|
||||
logger=logger,
|
||||
num_sanity_val_steps=0,
|
||||
|
||||
Reference in New Issue
Block a user