add ref + fix linting

Author: pri1311

test/test_ops.py

@@ -1767,28 +1767,22 @@ def test_is_leaf_node(self, dim, p, block_size, inplace):
 
 class TestDiceLoss:
     def get_reduction_method(self, reduction):
-        return {
-            "sum": torch.sum,
-            "mean": torch.mean,
-            "none": None
-        }[reduction]
+        return {"sum": torch.sum, "mean": torch.mean, "none": None}[reduction]
 
     @pytest.mark.parametrize("device", cpu_and_gpu())
     def test_dice_loss(self, device):
-        input = torch.tensor([[[0.9409, 0.9524],
-                               [0.9220, 0.1094]],
-                              [[0.6802, 0.9570],
-                               [0.7949, 0.1499]],
-                              [[0.3298, 0.1094],
-                               [0.4401, 0.7536]],
-                              [[0.3340, 0.9563],
-                               [0.9895, 0.5045]]], device=device)
-        labels = torch.tensor([[[0, 1], [1, 0]],
-                               [[1, 0], [0, 1]],
-                               [[1, 1], [0, 0]],
-                               [[1, 0], [0, 1]]], device=device)
+        input_tensor = torch.tensor(
+            [
+                [[0.9409, 0.9524], [0.9220, 0.1094]],
+                [[0.6802, 0.9570], [0.7949, 0.1499]],
+                [[0.3298, 0.1094], [0.4401, 0.7536]],
+                [[0.3340, 0.9563], [0.9895, 0.5045]],
+            ],
+            device=device,
+        )
+        labels = torch.tensor([[[0, 1], [1, 0]], [[1, 0], [0, 1]], [[1, 1], [0, 0]], [[1, 0], [0, 1]]], device=device)
         expected = torch.tensor([0.4028, 0.6101, 0.5916, 0.6347], device=device)
-        torch.testing.assert_allclose(ops.dice_loss(input, labels, eps=0), expected)
+        torch.testing.assert_allclose(ops.dice_loss(input_tensor, labels, eps=0), expected)
 
     @pytest.mark.parametrize("shape", ((16, 2, 4, 4), (16, 4, 4, 4), (32, 2), (32, 2, 4, 4, 4)))
     @pytest.mark.parametrize("reduction", ["none", "mean", "sum"])
@@ -1819,7 +1813,9 @@ def test_gradcheck(self, device):
         shape = (16, 2, 4, 4)
         input_ones = torch.ones(shape, device=device, requires_grad=True)
         label_zeros = torch.zeros(shape, device=device, requires_grad=True)
-        assert gradcheck(ops.dice_loss, (input_ones, label_zeros), eps=1e-2, atol=1e-2, raise_exception=True, fast_mode=True)
+        assert gradcheck(
+            ops.dice_loss, (input_ones, label_zeros), eps=1e-2, atol=1e-2, raise_exception=True, fast_mode=True
+        )
 
 
 if __name__ == "__main__":

torchvision/ops/dice_loss.py

@@ -4,6 +4,7 @@
 from ..utils import _log_api_usage_once
 
 
+# Implementation adapted from https://kornia.readthedocs.io/en/latest/_modules/kornia/losses/dice.html#dice_loss
 def dice_loss(inputs: torch.Tensor, targets: torch.Tensor, reduction: str = "none", eps: float = 1e-7) -> torch.Tensor:
     r"""Criterion that computes Sørensen-Dice Coefficient loss.
 
@@ -24,8 +25,8 @@ def dice_loss(inputs: torch.Tensor, targets: torch.Tensor, reduction: str = "non
         \text{loss}(x, class) = 1 - \text{Dice}(x, class)
 
     Args:
-        inputs: (Tensor): A float tensor with rank >= 2 and shape (B, C, N1, .... NK)
-                where B is the Batch Size and C is the number of classes.
+        inputs: (Tensor): A float tensor with rank >= 2 and shape (B, num_classes, N1, .... NK)
+                where B is the Batch Size and num_classes is the number of classes.
                 The predictions for each example.
         targets: (Tensor): A one-hot tensor with the same shape as inputs.
                 The first dimension is the batch size and the second dimension is the