<p dir="auto">PyTorch checks if tensors have been modified in-place to ensure correctness of the backward operation, with <code class="notranslate">ctx.save_for_backward</code> and <code class="notranslate">ctx.saved_tensors</code> in the <code class="notranslate">forward</code> and <code class="notranslate">backward</code> methods of subclasses of <code class="notranslate">torch.autograd.Function</code>. This check does not appear to function when the saved tensor is not the 'owner' of its storage, e.g. when it is the result of a <code class="notranslate">view</code> or <code class="notranslate">transpose</code> operation.</p>
<p dir="auto">To reproduce:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import torch

class f(torch.autograd.Function):
    @staticmethod
    def forward(ctx, tensor1):
        # make sure we don't share storage with our input
        # (just to be clear that that's not what's going wrong)
        tensor2 = tensor1.clone()
        # tensor3 shares storage with tensor2
        # tensor3 = tensor2.view(-1) etc. would also demonstrate the problem
        tensor3 = tensor2.t()
        ctx.save_for_backward(tensor3)
        return tensor3
    @staticmethod
    def backward(ctx, grad):
        # trigger correctness check
        _ = ctx.saved_tensors
        return grad.t().clone()


x = torch.rand(2, 3, requires_grad=True)
y = f.apply(x)
y += 1
# Does not raise an error; this is the bug.
y.backward(torch.rand_like(y))


# In constrast, the following code correctly raises an error.
class ff(torch.autograd.Function):
    @staticmethod
    def forward(ctx, tensor1):
        tensor2 = tensor1.clone()
        ctx.save_for_backward(tensor2)
        return tensor2
    @staticmethod
    def backward(ctx, grad):
        _ = ctx.saved_tensors
        return grad.clone()

xx = torch.rand(2, 3, requires_grad=True)
yy = ff.apply(xx)
yy += 1
yy.backward(torch.rand_like(yy))"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">torch</span>

<span class="pl-k">class</span> <span class="pl-s1">f</span>(<span class="pl-s1">torch</span>.<span class="pl-s1">autograd</span>.<span class="pl-v">Function</span>):
    <span class="pl-en">@<span class="pl-s1">staticmethod</span></span>
    <span class="pl-k">def</span> <span class="pl-en">forward</span>(<span class="pl-s1">ctx</span>, <span class="pl-s1">tensor1</span>):
        <span class="pl-c"># make sure we don't share storage with our input</span>
        <span class="pl-c"># (just to be clear that that's not what's going wrong)</span>
        <span class="pl-s1">tensor2</span> <span class="pl-c1">=</span> <span class="pl-s1">tensor1</span>.<span class="pl-en">clone</span>()
        <span class="pl-c"># tensor3 shares storage with tensor2</span>
        <span class="pl-c"># tensor3 = tensor2.view(-1) etc. would also demonstrate the problem</span>
        <span class="pl-s1">tensor3</span> <span class="pl-c1">=</span> <span class="pl-s1">tensor2</span>.<span class="pl-en">t</span>()
        <span class="pl-s1">ctx</span>.<span class="pl-en">save_for_backward</span>(<span class="pl-s1">tensor3</span>)
        <span class="pl-k">return</span> <span class="pl-s1">tensor3</span>
    <span class="pl-en">@<span class="pl-s1">staticmethod</span></span>
    <span class="pl-k">def</span> <span class="pl-en">backward</span>(<span class="pl-s1">ctx</span>, <span class="pl-s1">grad</span>):
        <span class="pl-c"># trigger correctness check</span>
        <span class="pl-s1">_</span> <span class="pl-c1">=</span> <span class="pl-s1">ctx</span>.<span class="pl-s1">saved_tensors</span>
        <span class="pl-k">return</span> <span class="pl-s1">grad</span>.<span class="pl-en">t</span>().<span class="pl-en">clone</span>()


<span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">rand</span>(<span class="pl-c1">2</span>, <span class="pl-c1">3</span>, <span class="pl-s1">requires_grad</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-s1">f</span>.<span class="pl-en">apply</span>(<span class="pl-s1">x</span>)
<span class="pl-s1">y</span> <span class="pl-c1">+=</span> <span class="pl-c1">1</span>
<span class="pl-c"># Does not raise an error; this is the bug.</span>
<span class="pl-s1">y</span>.<span class="pl-en">backward</span>(<span class="pl-s1">torch</span>.<span class="pl-en">rand_like</span>(<span class="pl-s1">y</span>))


<span class="pl-c"># In constrast, the following code correctly raises an error.</span>
<span class="pl-k">class</span> <span class="pl-s1">ff</span>(<span class="pl-s1">torch</span>.<span class="pl-s1">autograd</span>.<span class="pl-v">Function</span>):
    <span class="pl-en">@<span class="pl-s1">staticmethod</span></span>
    <span class="pl-k">def</span> <span class="pl-en">forward</span>(<span class="pl-s1">ctx</span>, <span class="pl-s1">tensor1</span>):
        <span class="pl-s1">tensor2</span> <span class="pl-c1">=</span> <span class="pl-s1">tensor1</span>.<span class="pl-en">clone</span>()
        <span class="pl-s1">ctx</span>.<span class="pl-en">save_for_backward</span>(<span class="pl-s1">tensor2</span>)
        <span class="pl-k">return</span> <span class="pl-s1">tensor2</span>
    <span class="pl-en">@<span class="pl-s1">staticmethod</span></span>
    <span class="pl-k">def</span> <span class="pl-en">backward</span>(<span class="pl-s1">ctx</span>, <span class="pl-s1">grad</span>):
        <span class="pl-s1">_</span> <span class="pl-c1">=</span> <span class="pl-s1">ctx</span>.<span class="pl-s1">saved_tensors</span>
        <span class="pl-k">return</span> <span class="pl-s1">grad</span>.<span class="pl-en">clone</span>()

<span class="pl-s1">xx</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">rand</span>(<span class="pl-c1">2</span>, <span class="pl-c1">3</span>, <span class="pl-s1">requires_grad</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)
<span class="pl-s1">yy</span> <span class="pl-c1">=</span> <span class="pl-s1">ff</span>.<span class="pl-en">apply</span>(<span class="pl-s1">xx</span>)
<span class="pl-s1">yy</span> <span class="pl-c1">+=</span> <span class="pl-c1">1</span>
<span class="pl-s1">yy</span>.<span class="pl-en">backward</span>(<span class="pl-s1">torch</span>.<span class="pl-en">rand_like</span>(<span class="pl-s1">yy</span>))</pre></div>
<p dir="auto">Expected behaviour is for an exception to be raised on the line <code class="notranslate">y.backward(torch.rand_like(y))</code>, in the same manner as an exception is raised on the line <code class="notranslate">yy.backward(torch.rand_like(yy))</code>.</p>
<ul dir="auto">
<li>PyTorch Version (e.g., 1.0): 1.2</li>
<li>OS (e.g., Linux): Linux</li>
<li>How you installed PyTorch (<code class="notranslate">conda</code>, <code class="notranslate">pip</code>, source): conda</li>
<li>Python version: 3.7.4</li>
</ul>
<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></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>