Apply suggestions from code review

Lightning-AI · Jun 18, 2020 · f50d991 · f50d991
1 parent 287361a
commit f50d991
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 16 deletions.
diff --git a/pytorch_lightning/metrics/classification.py b/pytorch_lightning/metrics/classification.py
@@ -16,7 +16,7 @@
     multiclass_roc,
     multiclass_precision_recall_curve,
     dice_score,
-    iou
+    iou,
 )
 from pytorch_lightning.metrics.metric import TensorMetric, TensorCollectionMetric
 

diff --git a/tests/metrics/functional/test_classification.py b/tests/metrics/functional/test_classification.py
@@ -367,24 +367,21 @@ def test_dice_score(pred, target, expected):
     assert score == expected
 
 
-@pytest.mark.parametrize(['target', 'pred', 'half_ones', 'reduction', 'remove_bg', 'expected'], [
-    pytest.param((torch.arange(120) % 3).view(-1, 1), (torch.arange(120) % 3).view(-1, 1),
-                 False, 'none', False, torch.Tensor([1, 1, 1])),
-    pytest.param((torch.arange(120) % 3).view(-1, 1), (torch.arange(120) % 3).view(-1, 1),
-                 False, 'elementwise_mean', False, torch.Tensor([1])),
-    pytest.param((torch.arange(120) % 3).view(-1, 1), (torch.arange(120) % 3).view(-1, 1),
-                 False, 'none', True, torch.Tensor([1, 1])),
-    pytest.param((torch.arange(120) % 3).view(-1, 1), (torch.arange(120) % 3).view(-1, 1),
-                 True, 'none', False, torch.Tensor([0.5, 0.5, 0.5])),
-    pytest.param((torch.arange(120) % 3).view(-1, 1), (torch.arange(120) % 3).view(-1, 1),
-                 True, 'elementwise_mean', False, torch.Tensor([0.5])),
-    pytest.param((torch.arange(120) % 3).view(-1, 1), (torch.arange(120) % 3).view(-1, 1),
-                 True, 'none', True, torch.Tensor([0.5, 0.5])),
+@pytest.mark.parametrize(['half_ones', 'reduction', 'remove_bg', 'expected'], [
+    pytest.param(False, 'none', False, torch.Tensor([1, 1, 1])),
+    pytest.param(False, 'elementwise_mean', False, torch.Tensor([1])),
+    pytest.param(False, 'none', True, torch.Tensor([1, 1])),
+    pytest.param(True, 'none', False, torch.Tensor([0.5, 0.5, 0.5])),
+    pytest.param(True, 'elementwise_mean', False, torch.Tensor([0.5])),
+    pytest.param(True, 'none', True, torch.Tensor([0.5, 0.5])),
 ])
-def test_iou(target, pred, half_ones, reduction, remove_bg, expected):
+def test_iou(half_ones, reduction, remove_bg, expected):
+    pred = (torch.arange(120) % 3).view(-1, 1)
+    target = (torch.arange(120) % 3).view(-1, 1)
     if half_ones:
         pred[:60] = 1
-    assert torch.all(torch.eq(iou(pred, target, remove_bg=remove_bg, reduction=reduction), expected))
+    iou_val = iou(pred, target, remove_bg=remove_bg, reduction=reduction)
+    assert torch.allclose(iou_val, expected, atol=1e-9)
 
 
 # example data taken from