Repository: kornia/kornia

Issue:
Gradient computation failed for distance_transform
### Describe the bug

Using `kornia.contrib.distance_transform` before computing an MSE loss would raise RuntimeError for gradient computation. 
The error message looks like this:

```
{...}/site-packages/torch/autograd/__init__.py:145: UserWarning: Error detected in ReplicationPad2DBackward. Traceback of forward call that caused the error:
  File "distance_transform_loss.py", line 104, in <module>
    a_dist = distance_transform(a)
  File "{...}/site-packages/kornia/contrib/distance_transform.py", line 59, in distance_transform
    cdt = filter2d(boundary, kernel, border_type='replicate')
  File "{...}/site-packages/kornia/filters/filter.py", line 114, in filter2d
    input = F.pad(input, padding_shape, mode=border_type)
  File "{...}/site-packages/torch/nn/functional.py", line 4019, in _pad
    return torch._C._nn.replication_pad2d(input, pad)
 (Triggered internally at  /opt/conda/conda-bld/pytorch_1616554788289/work/torch/csrc/autograd/python_anomaly_mode.cpp:104.)
  Variable._execution_engine.run_backward(
Traceback (most recent call last):
  File "distance_transform_loss.py", line 110, in <module>
    loss.backward()
  File "{...}/site-packages/torch/tensor.py", line 245, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "{...}/site-packages/torch/autograd/__init__.py", line 145, in backward
    Variable._execution_engine.run_backward(
RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [12, 1, 384, 384]], which is output 0 of IndexPutBackward, is at version 384; expected version 383 instead. Hint: the backtrace further above shows the operation that failed to compute its gradient. The variable in question was changed in there or anywhere later. Good luck!
```

A minimum reproducible script is provided below.


### Reproduction steps

```bash
import torch
torch.autograd.set_detect_anomaly(True)
import torch.nn.functional as F

from kornia.contrib import distance_transform

a = torch.rand((12, 1, 384, 384)).to(torch.float32)
b = torch.rand((12, 1, 384, 384)).to(torch.float32)

layer = nn.Conv2d(1, 1, (3, 3), (1, 1), (1, 1))
a = layer(a)

a_dist = distance_transform(a)
b_dist = distance_transform(b)
loss = F.mse_loss(a_dist, b_dist)

loss.backward()
```


### Expected behavior

Gradient back-propagated successfully and there should not be any console outputs

### Environment

```shell
wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
```
- PyTorch Version (e.g., 1.0): 1.8.1
- OS (e.g., Linux): Ubuntu 18.04.5 LTS (x86_64)
- How you installed PyTorch (`conda`, `pip`, source): conda
- Build command you used (if compiling from source): None
- Python version: 3.8.10
- CUDA/cuDNN version: 10.1.243/7.6.5
- GPU models and configuration: Tesla V100-SXM2-16GB
- Any other relevant information:
```


### Additional context

_No response_


Discussion:
@dichen-cd thanks to report. A quick guess is due to the inplace addition here
https://github.com/kornia/kornia/blob/master/kornia/contrib/distance_transform.py#L70
Hi @edgarriba . Thanks for the quick reply.

I tried to replace the inplace addtion `out += (offset + cdt) * mask` with `out = out + (offset + cdt) * mask`, the problem still remains.
> Hi @edgarriba . Thanks for the quick reply.
> 
> 
> 
> I tried to replace the inplace addtion `out += (offset + cdt) * mask` with `out = out + (offset + cdt) * mask`, the problem still remains.

How about the next line? Considering using torch.where instead?
Hi @shijianjian. Thank you for the suggestion!

I changed `boundary[mask == 1] = 1` to `boundary = torch.where(mask == 1, torch.ones_like(boundary), boundary)`.
The problem is now fixed.
> Hi @shijianjian. Thank you for the suggestion!
> 
> 
> 
> I changed `boundary[mask == 1] = 1` to `boundary = torch.where(mask == 1, torch.ones_like(boundary), boundary)`.
> 
> The problem is now fixed.

Brilliant. Do you mind wrapping up and sending a PR?

Resolving patch:
diff --git a/kornia/contrib/distance_transform.py b/kornia/contrib/distance_transform.py
--- a/kornia/contrib/distance_transform.py
+++ b/kornia/contrib/distance_transform.py
@@ -54,6 +54,7 @@ def distance_transform(
 
     # It is possible to avoid cloning the image if boundary = image, but this would require modifying the image tensor.
     boundary = image.clone()
+    signal_ones = torch.ones_like(boundary)
 
     for i in range(n_iters):
         cdt = filter2d(boundary, kernel, border_type='replicate')
@@ -68,7 +69,7 @@ def distance_transform(
 
         offset: int = i * kernel_size // 2
         out += (offset + cdt) * mask
-        boundary[mask == 1] = 1
+        boundary = torch.where(mask == 1, signal_ones, boundary)
 
     return out
 


Test patch:
diff --git a/test/test_contrib.py b/test/test_contrib.py
--- a/test/test_contrib.py
+++ b/test/test_contrib.py
@@ -436,6 +436,16 @@ def test_gradcheck(self, device, dtype):
         input1 = torch.ones(B, C, H, W, device=device, dtype=dtype, requires_grad=True)
         assert gradcheck(kornia.contrib.distance_transform, (input1), raise_exception=True)
 
+    def test_loss_grad(self, device, dtype):
+        B, C, H, W = 1, 1, 32, 32
+        input1 = torch.rand(B, C, H, W, device=device, dtype=dtype, requires_grad=True)
+        input2 = torch.rand(B, C, H, W, device=device, dtype=dtype, requires_grad=True)
+        tiny_module = torch.nn.Conv2d(1, 1, (3, 3), (1, 1), (1, 1)).to(device=device, dtype=dtype)
+        input1 = kornia.contrib.distance_transform(tiny_module(input1))
+        input2 = kornia.contrib.distance_transform(input2)
+        loss = torch.nn.functional.mse_loss(input1, input2)
+        loss.backward()
+
 
 class TestHistMatch:
     def test_interp(self, device, dtype):