From c1b3780de3598134abbe066770ee0c485a6322bb Mon Sep 17 00:00:00 2001
From: "Daniel J. Hofmann" <daniel@trvx.org>
Date: Tue, 2 Oct 2018 12:47:38 +0200
Subject: [PATCH] Fixes missing DataParallel in `rs train`

---
 robosat/tools/train.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/robosat/tools/train.py b/robosat/tools/train.py
index c8e8fee3..f6b74079 100644
--- 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