Support training on CPU (open-mmlab#720)

Yshuo-Li · Feb 10, 2022 · aa1bf8d · aa1bf8d
1 parent 6ac343d
commit aa1bf8d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mmedit/apis/train.py b/mmedit/apis/train.py
@@ -254,8 +254,8 @@ def _non_dist_train(model,
 
     data_loaders = [build_dataloader(ds, **train_loader_cfg) for ds in dataset]
 
-    # put model on gpus
-    model = MMDataParallel(model, device_ids=range(cfg.gpus)).cuda()
+    # put model on gpus/cpus
+    model = MMDataParallel(model, device_ids=range(cfg.gpus))
 
     # build runner
     optimizer = build_optimizers(model, cfg.optimizers)