fix mmcv version

83363a9a · Tao Kong · GitHub · 1e79dac4 · 83363a9a
Unverified Commit 83363a9a authored Apr 01, 2020 by Tao Kong Committed by GitHub Apr 01, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

mmdet/apis/train.py mmdet/apis/train.py +4 -1

No files found.
--- a/mmdet/apis/train.py
+++ b/mmdet/apis/train.py
@@ -205,7 +205,10 @@ def _dist_train(model,
        for ds in dataset
    ]
    # put model on gpus
-    model = MMDistributedDataParallel(model.cuda())
+    # model = MMDistributedDataParallel(model.cuda())
+    model = MMDistributedDataParallel(model.cuda(), 
+                                      device_ids=[torch.cuda.current_device()], 
+                                      broadcast_buffers=False)

    # build runner
    optimizer = build_optimizer(model, cfg.optimizer)