- made metrics friendly for multi-gpu model parallel training

inferno/extensions/metrics/base.py

@@ -5,5 +5,14 @@ class Metric(object):
     def forward(self, *args, **kwargs):
         raise NotImplementedError
 
-    def __call__(self, *args, **kwargs):
-        return self.forward(*args, **kwargs)
+    def __call__(self, prediction, target, **kwargs):
+        # Make sure prediction and target live on the same device.
+        # If they don't, move target to the right device.
+        if not prediction.is_cuda:
+            # Move to CPU
+            target = target.cpu()
+        else:
+            # Find device to move to
+            device_ordinal = prediction.get_device()
+            target = target.cuda(device_ordinal)
+        return self.forward(prediction, target, **kwargs)

inferno/extensions/metrics/categorical.py

@@ -118,7 +118,8 @@ def forward(self, prediction, target):
             dont_ignore_class = list(range(num_classes))
             dont_ignore_class.pop(ignore_class)
             if classwise_iou.is_cuda:
-                dont_ignore_class = torch.cuda.LongTensor(dont_ignore_class)
+                dont_ignore_class = \
+                    torch.LongTensor(dont_ignore_class).cuda(classwise_iou.get_device())
             else:
                 dont_ignore_class = torch.LongTensor(dont_ignore_class)
             iou = classwise_iou[dont_ignore_class].mean()