Fixes missing DataParallel in rs train

mapbox · Oct 2, 2018 · c1b3780 · c1b3780
1 parent 344c0bb
commit c1b3780
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/robosat/tools/train.py b/robosat/tools/train.py
@@ -63,7 +63,9 @@ def main(args):
     os.makedirs(model["common"]["checkpoint"], exist_ok=True)
 
     num_classes = len(dataset["common"]["classes"])
-    net = FPNSegmentation(num_classes).to(device)
+    net = FPNSegmentation(num_classes)
+    net = DataParallel(net)
+    net = net.to(device)
 
     if model["common"]["cuda"]:
         torch.backends.cudnn.benchmark = True