<p dir="auto">If the output of a python autograd function is a view and gets modified in-place with a torch.no_grad block later, the python autograd function's backward function doesn't get called and instead an autograd backward pass (through the forward function) runs. This results in incorrect behavior, especially when the forward pass contains non-differentiable things.</p>
<p dir="auto">Expected behavior: we should either throw an error message about something or compute the correct gradients. Neither of those happen right now.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import torch
from torch.autograd import Function

class Mul3ThenFlatten(Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        z = x.detach().numpy() * 3
        # If I add a .contiguous() to the end of this statement, other things happen.
        return torch.tensor(z).view(-1)

    @staticmethod
    def backward(ctx, dx):
        x = ctx.saved_tensors[0]
        print(&quot;Some side effect&quot;)
        return dx.view(x.shape) * 3

x = torch.randn(2, 2, requires_grad=True)
output = Mul3ThenFlatten.apply(x)

# Modify the counter in no_grad mode.
with torch.no_grad():
    output.fill_(1)

# Doesn't print &quot;Some side effect&quot; :/
output.sum().backward()

x.grad  # is None, should really be something or at least an assert"><pre class="notranslate"><code class="notranslate">import torch
from torch.autograd import Function

class Mul3ThenFlatten(Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        z = x.detach().numpy() * 3
        # If I add a .contiguous() to the end of this statement, other things happen.
        return torch.tensor(z).view(-1)

    @staticmethod
    def backward(ctx, dx):
        x = ctx.saved_tensors[0]
        print("Some side effect")
        return dx.view(x.shape) * 3

x = torch.randn(2, 2, requires_grad=True)
output = Mul3ThenFlatten.apply(x)

# Modify the counter in no_grad mode.
with torch.no_grad():
    output.fill_(1)

# Doesn't print "Some side effect" :/
output.sum().backward()

x.grad  # is None, should really be something or at least an assert
</code></pre></div>
<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/colesbury/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/colesbury">@colesbury</a></p>
<h2 dir="auto"><g-emoji class="g-emoji" alias="bug" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/1f41b.png">🐛</g-emoji> Bug</h2>
<p dir="auto">Gradients aren't backpropagated when all three conditions are met:</p>
<ol dir="auto">
<li>Use <code class="notranslate">torch.utils.checkpoint</code>.</li>
<li>Place <code class="notranslate">nn.InstanceNorm?d</code> right before any in-place module such as <code class="notranslate">nn.ReLU(inplace=True)</code>.</li>
<li>The checkpoint returns the result of the in-place module.</li>
</ol>
<h2 dir="auto">To Reproduce</h2>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import torch
from torch import nn
from torch.utils.checkpoint import checkpoint

x = torch.ones(3, 2, 1, requires_grad=True)

model = nn.Sequential(
    nn.InstanceNorm1d(2),
    nn.ReLU(inplace=True),
)
y = checkpoint(model, x)

y.norm().backward()

print(x.grad)
# Output: None
# Expected: tensor([[[0.],[0.]], ...])"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">torch</span>
<span class="pl-k">from</span> <span class="pl-s1">torch</span> <span class="pl-k">import</span> <span class="pl-s1">nn</span>
<span class="pl-k">from</span> <span class="pl-s1">torch</span>.<span class="pl-s1">utils</span>.<span class="pl-s1">checkpoint</span> <span class="pl-k">import</span> <span class="pl-s1">checkpoint</span>

<span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">ones</span>(<span class="pl-c1">3</span>, <span class="pl-c1">2</span>, <span class="pl-c1">1</span>, <span class="pl-s1">requires_grad</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)

<span class="pl-s1">model</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">Sequential</span>(
    <span class="pl-s1">nn</span>.<span class="pl-v">InstanceNorm1d</span>(<span class="pl-c1">2</span>),
    <span class="pl-s1">nn</span>.<span class="pl-v">ReLU</span>(<span class="pl-s1">inplace</span><span class="pl-c1">=</span><span class="pl-c1">True</span>),
)
<span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-en">checkpoint</span>(<span class="pl-s1">model</span>, <span class="pl-s1">x</span>)

<span class="pl-s1">y</span>.<span class="pl-en">norm</span>().<span class="pl-en">backward</span>()

<span class="pl-en">print</span>(<span class="pl-s1">x</span>.<span class="pl-s1">grad</span>)
<span class="pl-c"># Output: None</span>
<span class="pl-c"># Expected: tensor([[[0.],[0.]], ...])</span></pre></div>
<p dir="auto"><code class="notranslate">x</code> the input tensor doesn't receive gradients. <code class="notranslate">x.grad</code> is still <code class="notranslate">None</code>, which is not expected.</p>
<h2 dir="auto">Expected behavior</h2>
<p dir="auto"><code class="notranslate">x.grad</code> should be <code class="notranslate">torch.zeros(3, 2, 1)</code>. If we remove any of the conditions, the gradients are backpropagated correctly:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="# No in-place operation
model = nn.Sequential(
    nn.InstanceNorm1d(2),
    nn.ReLU(inplace=False),
)
y = checkpoint(model, x)

# No checkpoint
model = nn.Sequential(
    nn.InstanceNorm1d(2),
    nn.ReLU(inplace=True),
)
y = model(x)

# Not at the end of checkpoint
model = nn.Sequential(
    nn.InstanceNorm1d(2),
    nn.ReLU(inplace=True),
    nn.Conv1d(2, 2, 1),
)
y = checkpoint(model, x)

# Other norm
model = nn.Sequential(
    nn.BatchNorm1d(2),
    nn.ReLU(inplace=True),
)
y = checkpoint(model, x)"><pre class="notranslate"><span class="pl-c"># No in-place operation</span>
<span class="pl-s1">model</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">Sequential</span>(
    <span class="pl-s1">nn</span>.<span class="pl-v">InstanceNorm1d</span>(<span class="pl-c1">2</span>),
    <span class="pl-s1">nn</span>.<span class="pl-v">ReLU</span>(<span class="pl-s1">inplace</span><span class="pl-c1">=</span><span class="pl-c1">False</span>),
)
<span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-en">checkpoint</span>(<span class="pl-s1">model</span>, <span class="pl-s1">x</span>)

<span class="pl-c"># No checkpoint</span>
<span class="pl-s1">model</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">Sequential</span>(
    <span class="pl-s1">nn</span>.<span class="pl-v">InstanceNorm1d</span>(<span class="pl-c1">2</span>),
    <span class="pl-s1">nn</span>.<span class="pl-v">ReLU</span>(<span class="pl-s1">inplace</span><span class="pl-c1">=</span><span class="pl-c1">True</span>),
)
<span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-en">model</span>(<span class="pl-s1">x</span>)

<span class="pl-c"># Not at the end of checkpoint</span>
<span class="pl-s1">model</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">Sequential</span>(
    <span class="pl-s1">nn</span>.<span class="pl-v">InstanceNorm1d</span>(<span class="pl-c1">2</span>),
    <span class="pl-s1">nn</span>.<span class="pl-v">ReLU</span>(<span class="pl-s1">inplace</span><span class="pl-c1">=</span><span class="pl-c1">True</span>),
    <span class="pl-s1">nn</span>.<span class="pl-v">Conv1d</span>(<span class="pl-c1">2</span>, <span class="pl-c1">2</span>, <span class="pl-c1">1</span>),
)
<span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-en">checkpoint</span>(<span class="pl-s1">model</span>, <span class="pl-s1">x</span>)

<span class="pl-c"># Other norm</span>
<span class="pl-s1">model</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">Sequential</span>(
    <span class="pl-s1">nn</span>.<span class="pl-v">BatchNorm1d</span>(<span class="pl-c1">2</span>),
    <span class="pl-s1">nn</span>.<span class="pl-v">ReLU</span>(<span class="pl-s1">inplace</span><span class="pl-c1">=</span><span class="pl-c1">True</span>),
)
<span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-en">checkpoint</span>(<span class="pl-s1">model</span>, <span class="pl-s1">x</span>)</pre></div>
<h2 dir="auto">Environment</h2>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="Collecting environment information...
PyTorch version: 1.2.0
Is debug build: No
CUDA used to build PyTorch: 10.0.130

OS: Ubuntu 18.04.2 LTS
GCC version: (Ubuntu 7.4.0-1ubuntu1~18.04.1) 7.4.0
CMake version: version 3.10.2

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: 10.0.130
GPU models and configuration: 
GPU 0: Tesla P40
GPU 1: Tesla P40
GPU 2: Tesla P40
GPU 3: Tesla P40
GPU 4: Tesla P40
GPU 5: Tesla P40
GPU 6: Tesla P40
GPU 7: Tesla P40

Nvidia driver version: 410.104
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.2

Versions of relevant libraries:
[pip3] numpy==1.16.4
[pip3] torch==1.2.0
[pip3] torchgpipe==0.0.3
[pip3] torchvision==0.4.0
[conda] blas                      1.0                         mkl  
[conda] mkl                       2019.4                      243  
[conda] mkl-include               2019.4                      243  
[conda] mkl-service               2.0.2            py36h7b6447c_0  
[conda] mkl_fft                   1.0.12           py36ha843d7b_0  
[conda] mkl_random                1.0.2            py36hd81dba3_0  
[conda] torch                     1.2.0                    pypi_0    pypi
[conda] torchgpipe                0.0.3                     dev_0    &lt;develop&gt;
[conda] torchvision               0.4.0                    pypi_0    pypi"><pre class="notranslate"><code class="notranslate">Collecting environment information...
PyTorch version: 1.2.0
Is debug build: No
CUDA used to build PyTorch: 10.0.130

OS: Ubuntu 18.04.2 LTS
GCC version: (Ubuntu 7.4.0-1ubuntu1~18.04.1) 7.4.0
CMake version: version 3.10.2

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: 10.0.130
GPU models and configuration: 
GPU 0: Tesla P40
GPU 1: Tesla P40
GPU 2: Tesla P40
GPU 3: Tesla P40
GPU 4: Tesla P40
GPU 5: Tesla P40
GPU 6: Tesla P40
GPU 7: Tesla P40

Nvidia driver version: 410.104
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.2

Versions of relevant libraries:
[pip3] numpy==1.16.4
[pip3] torch==1.2.0
[pip3] torchgpipe==0.0.3
[pip3] torchvision==0.4.0
[conda] blas                      1.0                         mkl  
[conda] mkl                       2019.4                      243  
[conda] mkl-include               2019.4                      243  
[conda] mkl-service               2.0.2            py36h7b6447c_0  
[conda] mkl_fft                   1.0.12           py36ha843d7b_0  
[conda] mkl_random                1.0.2            py36hd81dba3_0  
[conda] torch                     1.2.0                    pypi_0    pypi
[conda] torchgpipe                0.0.3                     dev_0    &lt;develop&gt;
[conda] torchvision               0.4.0                    pypi_0    pypi
</code></pre></div>
<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/ezyang/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/ezyang">@ezyang</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/ssnl/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/ssnl">@ssnl</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/albanD/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/albanD">@albanD</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/zou3519/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/zou3519">@zou3519</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/gqchen/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/gqchen">@gqchen</a></p>