bug fix for long label

Twihtronics · Nov 11, 2020 · d0da954 · d0da954
1 parent bb673eb
commit d0da954
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/train.py b/train.py
@@ -17,12 +17,12 @@
 def inference(net, data_label, use_aux):
     if use_aux:
         img, cls_label, seg_label = data_label
-        img, cls_label, seg_label = img.cuda(), cls_label.cuda(), seg_label.cuda()
+        img, cls_label, seg_label = img.cuda(), cls_label.long().cuda(), seg_label.long().cuda()
         cls_out, seg_out = net(img)
         return {'cls_out': cls_out, 'cls_label': cls_label, 'seg_out':seg_out, 'seg_label': seg_label}
     else:
         img, cls_label = data_label
-        img, cls_label = img.cuda(), cls_label.cuda()
+        img, cls_label = img.cuda(), cls_label.long().cuda()
         cls_out = net(img)
         return {'cls_out': cls_out, 'cls_label': cls_label}