<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">I want to use these two very useful functions at the same time: <code class="notranslate">torch.cuda.amp</code> and <code class="notranslate">torch.utils.checkpoint.checkpoint</code>, but I get this error:</p>
<div class="highlight highlight-source-shell notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="Traceback (most recent call last):
  File &quot;/home/lart/Coding/MyProj/utils/some_tools/test_amp_ck.py&quot;, line 72, in &lt;module&gt;
    scaler.scale(loss).backward()
  File &quot;/home/lart/miniconda3/envs/pt16end/lib/python3.8/site-packages/torch/tensor.py&quot;, line 185, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File &quot;/home/lart/miniconda3/envs/pt16end/lib/python3.8/site-packages/torch/autograd/__init__.py&quot;, line 125, in backward
    Variable._execution_engine.run_backward(
RuntimeError: Input type (torch.cuda.HalfTensor) and weight type (torch.cuda.FloatTensor) should be the same

Process finished with exit code 1"><pre class="notranslate">Traceback (most recent call last):
  File <span class="pl-s"><span class="pl-pds">"</span>/home/lart/Coding/MyProj/utils/some_tools/test_amp_ck.py<span class="pl-pds">"</span></span>, line 72, <span class="pl-k">in</span> <span class="pl-k">&lt;</span>module<span class="pl-k">&gt;</span>
    <span class="pl-en">scaler.scale(loss).backward</span>()
  File <span class="pl-s"><span class="pl-pds">"</span>/home/lart/miniconda3/envs/pt16end/lib/python3.8/site-packages/torch/tensor.py<span class="pl-pds">"</span></span>, line 185, <span class="pl-k">in</span> backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File <span class="pl-s"><span class="pl-pds">"</span>/home/lart/miniconda3/envs/pt16end/lib/python3.8/site-packages/torch/autograd/__init__.py<span class="pl-pds">"</span></span>, line 125, <span class="pl-k">in</span> backward
    Variable._execution_engine.run_backward(
RuntimeError: Input <span class="pl-c1">type</span> (torch.cuda.HalfTensor) and weight <span class="pl-c1">type</span> (torch.cuda.FloatTensor) should be the same

Process finished with <span class="pl-c1">exit</span> code 1</pre></div>
<p dir="auto"><strong>BUT</strong> IF I USE the APEX lib from NVIDIA, It works well...</p>
<h2 dir="auto">To Reproduce</h2>

<p dir="auto">A demo:</p>
<p dir="auto">The shared part:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="# -*- coding: utf-8 -*-
# @Time    : 2020/11/2
# @Author  : Lart Pang
# @File    : test_amp_ck.py
# @Project : MyProj
# @GitHub  : https://github.com/lartpang
import apex
import torch
from torch.utils.checkpoint import checkpoint
from tqdm import tqdm


class MyNetWithCheckpoint(torch.nn.Module):
    def __init__(self):
        super(MyNetWithCheckpoint, self).__init__()
        self.dummy_tensor = torch.ones(1, dtype=torch.float32,
                                       requires_grad=True)
        self.in_conv = torch.nn.Conv2d(3, 10, 3, 1, 1)
        self.out_conv = torch.nn.Conv2d(10, 1, 1)

    def forward(self, x):
        x = checkpoint(self.ck_in_conv, x, self.dummy_tensor)
        x = checkpoint(self.ck_out_conv, x)
        return x

    def ck_in_conv(self, x, dummy_arg=None):
        assert dummy_arg is not None
        x = self.in_conv(x)
        return x

    def ck_out_conv(self, x):
        # assert dummy_arg is not None
        x = self.out_conv(x)
        return x


def cal_loss(pred, gt):
    pred = pred.sigmoid()
    loss = (pred - gt) ** 2
    return loss.mean()


mynet = MyNetWithCheckpoint().cuda()

optimizer = torch.optim.SGD(mynet.parameters(), lr=0.001)
USE_AMP = True
ITER_NUM = 10

in_data = torch.randn(4, 3, 20, 20).cuda()
in_gt = torch.randint(low=0, high=2, size=(4, 3, 20, 20)).float().cuda()"><pre class="notranslate"><span class="pl-c"># -*- coding: utf-8 -*-</span>
<span class="pl-c"># @Time    : 2020/11/2</span>
<span class="pl-c"># @Author  : Lart Pang</span>
<span class="pl-c"># @File    : test_amp_ck.py</span>
<span class="pl-c"># @Project : MyProj</span>
<span class="pl-c"># @GitHub  : https://github.com/lartpang</span>
<span class="pl-k">import</span> <span class="pl-s1">apex</span>
<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-s1">utils</span>.<span class="pl-s1">checkpoint</span> <span class="pl-k">import</span> <span class="pl-s1">checkpoint</span>
<span class="pl-k">from</span> <span class="pl-s1">tqdm</span> <span class="pl-k">import</span> <span class="pl-s1">tqdm</span>


<span class="pl-k">class</span> <span class="pl-v">MyNetWithCheckpoint</span>(<span class="pl-s1">torch</span>.<span class="pl-s1">nn</span>.<span class="pl-v">Module</span>):
    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>):
        <span class="pl-en">super</span>(<span class="pl-v">MyNetWithCheckpoint</span>, <span class="pl-s1">self</span>).<span class="pl-en">__init__</span>()
        <span class="pl-s1">self</span>.<span class="pl-s1">dummy_tensor</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">ones</span>(<span class="pl-c1">1</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">torch</span>.<span class="pl-s1">float32</span>,
                                       <span class="pl-s1">requires_grad</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)
        <span class="pl-s1">self</span>.<span class="pl-s1">in_conv</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-s1">nn</span>.<span class="pl-v">Conv2d</span>(<span class="pl-c1">3</span>, <span class="pl-c1">10</span>, <span class="pl-c1">3</span>, <span class="pl-c1">1</span>, <span class="pl-c1">1</span>)
        <span class="pl-s1">self</span>.<span class="pl-s1">out_conv</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-s1">nn</span>.<span class="pl-v">Conv2d</span>(<span class="pl-c1">10</span>, <span class="pl-c1">1</span>, <span class="pl-c1">1</span>)

    <span class="pl-k">def</span> <span class="pl-en">forward</span>(<span class="pl-s1">self</span>, <span class="pl-s1">x</span>):
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-en">checkpoint</span>(<span class="pl-s1">self</span>.<span class="pl-s1">ck_in_conv</span>, <span class="pl-s1">x</span>, <span class="pl-s1">self</span>.<span class="pl-s1">dummy_tensor</span>)
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-en">checkpoint</span>(<span class="pl-s1">self</span>.<span class="pl-s1">ck_out_conv</span>, <span class="pl-s1">x</span>)
        <span class="pl-k">return</span> <span class="pl-s1">x</span>

    <span class="pl-k">def</span> <span class="pl-en">ck_in_conv</span>(<span class="pl-s1">self</span>, <span class="pl-s1">x</span>, <span class="pl-s1">dummy_arg</span><span class="pl-c1">=</span><span class="pl-c1">None</span>):
        <span class="pl-k">assert</span> <span class="pl-s1">dummy_arg</span> <span class="pl-c1">is</span> <span class="pl-c1">not</span> <span class="pl-c1">None</span>
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">in_conv</span>(<span class="pl-s1">x</span>)
        <span class="pl-k">return</span> <span class="pl-s1">x</span>

    <span class="pl-k">def</span> <span class="pl-en">ck_out_conv</span>(<span class="pl-s1">self</span>, <span class="pl-s1">x</span>):
        <span class="pl-c"># assert dummy_arg is not None</span>
        <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">out_conv</span>(<span class="pl-s1">x</span>)
        <span class="pl-k">return</span> <span class="pl-s1">x</span>


<span class="pl-k">def</span> <span class="pl-en">cal_loss</span>(<span class="pl-s1">pred</span>, <span class="pl-s1">gt</span>):
    <span class="pl-s1">pred</span> <span class="pl-c1">=</span> <span class="pl-s1">pred</span>.<span class="pl-en">sigmoid</span>()
    <span class="pl-s1">loss</span> <span class="pl-c1">=</span> (<span class="pl-s1">pred</span> <span class="pl-c1">-</span> <span class="pl-s1">gt</span>) <span class="pl-c1">**</span> <span class="pl-c1">2</span>
    <span class="pl-k">return</span> <span class="pl-s1">loss</span>.<span class="pl-en">mean</span>()


<span class="pl-s1">mynet</span> <span class="pl-c1">=</span> <span class="pl-v">MyNetWithCheckpoint</span>().<span class="pl-en">cuda</span>()

<span class="pl-s1">optimizer</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-s1">optim</span>.<span class="pl-v">SGD</span>(<span class="pl-s1">mynet</span>.<span class="pl-en">parameters</span>(), <span class="pl-s1">lr</span><span class="pl-c1">=</span><span class="pl-c1">0.001</span>)
<span class="pl-v">USE_AMP</span> <span class="pl-c1">=</span> <span class="pl-c1">True</span>
<span class="pl-v">ITER_NUM</span> <span class="pl-c1">=</span> <span class="pl-c1">10</span>

<span class="pl-s1">in_data</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">randn</span>(<span class="pl-c1">4</span>, <span class="pl-c1">3</span>, <span class="pl-c1">20</span>, <span class="pl-c1">20</span>).<span class="pl-en">cuda</span>()
<span class="pl-s1">in_gt</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">randint</span>(<span class="pl-s1">low</span><span class="pl-c1">=</span><span class="pl-c1">0</span>, <span class="pl-s1">high</span><span class="pl-c1">=</span><span class="pl-c1">2</span>, <span class="pl-s1">size</span><span class="pl-c1">=</span>(<span class="pl-c1">4</span>, <span class="pl-c1">3</span>, <span class="pl-c1">20</span>, <span class="pl-c1">20</span>)).<span class="pl-en">float</span>().<span class="pl-en">cuda</span>()</pre></div>
<p dir="auto">When I use <code class="notranslate">torch.cuda.amp</code>:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="# with torch.cuda.amp
scaler = torch.cuda.amp.GradScaler(USE_AMP)

for _ in tqdm(range(ITER_NUM), total=ITER_NUM):
    with torch.cuda.amp.autocast(USE_AMP):
        pred = mynet(in_data)
    loss = cal_loss(pred=pred, gt=in_gt)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()"><pre class="notranslate"><span class="pl-c"># with torch.cuda.amp</span>
<span class="pl-s1">scaler</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-s1">cuda</span>.<span class="pl-s1">amp</span>.<span class="pl-v">GradScaler</span>(<span class="pl-v">USE_AMP</span>)

<span class="pl-k">for</span> <span class="pl-s1">_</span> <span class="pl-c1">in</span> <span class="pl-en">tqdm</span>(<span class="pl-en">range</span>(<span class="pl-v">ITER_NUM</span>), <span class="pl-s1">total</span><span class="pl-c1">=</span><span class="pl-v">ITER_NUM</span>):
    <span class="pl-k">with</span> <span class="pl-s1">torch</span>.<span class="pl-s1">cuda</span>.<span class="pl-s1">amp</span>.<span class="pl-en">autocast</span>(<span class="pl-v">USE_AMP</span>):
        <span class="pl-s1">pred</span> <span class="pl-c1">=</span> <span class="pl-en">mynet</span>(<span class="pl-s1">in_data</span>)
    <span class="pl-s1">loss</span> <span class="pl-c1">=</span> <span class="pl-en">cal_loss</span>(<span class="pl-s1">pred</span><span class="pl-c1">=</span><span class="pl-s1">pred</span>, <span class="pl-s1">gt</span><span class="pl-c1">=</span><span class="pl-s1">in_gt</span>)
    <span class="pl-s1">scaler</span>.<span class="pl-en">scale</span>(<span class="pl-s1">loss</span>).<span class="pl-en">backward</span>()
    <span class="pl-s1">scaler</span>.<span class="pl-en">step</span>(<span class="pl-s1">optimizer</span>)
    <span class="pl-s1">scaler</span>.<span class="pl-en">update</span>()
    <span class="pl-s1">optimizer</span>.<span class="pl-en">zero_grad</span>()</pre></div>
<p dir="auto">It will output:</p>
<div class="highlight highlight-source-shell notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="Traceback (most recent call last):
  File &quot;/home/lart/Coding/MyProj/utils/some_tools/test_amp_ck.py&quot;, line 72, in &lt;module&gt;
    scaler.scale(loss).backward()
  File &quot;/home/lart/miniconda3/envs/pt16end/lib/python3.8/site-packages/torch/tensor.py&quot;, line 185, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File &quot;/home/lart/miniconda3/envs/pt16end/lib/python3.8/site-packages/torch/autograd/__init__.py&quot;, line 125, in backward
    Variable._execution_engine.run_backward(
RuntimeError: Input type (torch.cuda.HalfTensor) and weight type (torch.cuda.FloatTensor) should be the same"><pre class="notranslate">Traceback (most recent call last):
  File <span class="pl-s"><span class="pl-pds">"</span>/home/lart/Coding/MyProj/utils/some_tools/test_amp_ck.py<span class="pl-pds">"</span></span>, line 72, <span class="pl-k">in</span> <span class="pl-k">&lt;</span>module<span class="pl-k">&gt;</span>
    <span class="pl-en">scaler.scale(loss).backward</span>()
  File <span class="pl-s"><span class="pl-pds">"</span>/home/lart/miniconda3/envs/pt16end/lib/python3.8/site-packages/torch/tensor.py<span class="pl-pds">"</span></span>, line 185, <span class="pl-k">in</span> backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File <span class="pl-s"><span class="pl-pds">"</span>/home/lart/miniconda3/envs/pt16end/lib/python3.8/site-packages/torch/autograd/__init__.py<span class="pl-pds">"</span></span>, line 125, <span class="pl-k">in</span> backward
    Variable._execution_engine.run_backward(
RuntimeError: Input <span class="pl-c1">type</span> (torch.cuda.HalfTensor) and weight <span class="pl-c1">type</span> (torch.cuda.FloatTensor) should be the same</pre></div>
<p dir="auto">When I use APEX from NVIDIA:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="# with apex.amp
mynet, optimizer = apex.amp.initialize(mynet, optimizer, opt_level=&quot;O1&quot;)

for _ in tqdm(range(ITER_NUM), total=ITER_NUM):
    pred = mynet(in_data)
    loss = cal_loss(pred=pred, gt=in_gt)
    if USE_AMP:
        with apex.amp.scale_loss(loss, optimizer) as scaled_loss:
            scaled_loss.backward()
    else:
        loss.backward()
    optimizer.step()
    optimizer.zero_grad()"><pre class="notranslate"><span class="pl-c"># with apex.amp</span>
<span class="pl-s1">mynet</span>, <span class="pl-s1">optimizer</span> <span class="pl-c1">=</span> <span class="pl-s1">apex</span>.<span class="pl-s1">amp</span>.<span class="pl-en">initialize</span>(<span class="pl-s1">mynet</span>, <span class="pl-s1">optimizer</span>, <span class="pl-s1">opt_level</span><span class="pl-c1">=</span><span class="pl-s">"O1"</span>)

<span class="pl-k">for</span> <span class="pl-s1">_</span> <span class="pl-c1">in</span> <span class="pl-en">tqdm</span>(<span class="pl-en">range</span>(<span class="pl-v">ITER_NUM</span>), <span class="pl-s1">total</span><span class="pl-c1">=</span><span class="pl-v">ITER_NUM</span>):
    <span class="pl-s1">pred</span> <span class="pl-c1">=</span> <span class="pl-en">mynet</span>(<span class="pl-s1">in_data</span>)
    <span class="pl-s1">loss</span> <span class="pl-c1">=</span> <span class="pl-en">cal_loss</span>(<span class="pl-s1">pred</span><span class="pl-c1">=</span><span class="pl-s1">pred</span>, <span class="pl-s1">gt</span><span class="pl-c1">=</span><span class="pl-s1">in_gt</span>)
    <span class="pl-k">if</span> <span class="pl-v">USE_AMP</span>:
        <span class="pl-k">with</span> <span class="pl-s1">apex</span>.<span class="pl-s1">amp</span>.<span class="pl-en">scale_loss</span>(<span class="pl-s1">loss</span>, <span class="pl-s1">optimizer</span>) <span class="pl-k">as</span> <span class="pl-s1">scaled_loss</span>:
            <span class="pl-s1">scaled_loss</span>.<span class="pl-en">backward</span>()
    <span class="pl-k">else</span>:
        <span class="pl-s1">loss</span>.<span class="pl-en">backward</span>()
    <span class="pl-s1">optimizer</span>.<span class="pl-en">step</span>()
    <span class="pl-s1">optimizer</span>.<span class="pl-en">zero_grad</span>()</pre></div>
<p dir="auto">It will output:</p>
<div class="highlight highlight-source-shell notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="
Selected optimization level O1:  Insert automatic casts around Pytorch functions and Tensor methods.

Defaults for this optimization level are:
enabled                : True
opt_level              : O1
cast_model_type        : None
patch_torch_functions  : True
keep_batchnorm_fp32    : None
master_weights         : None
loss_scale             : dynamic
Processing user overrides (additional kwargs that are not None)...
After processing overrides, optimization options are:
enabled                : True
opt_level              : O1
cast_model_type        : None
patch_torch_functions  : True
keep_batchnorm_fp32    : None
master_weights         : None
loss_scale             : dynamic
100%|██████████| 10/10 [00:00&lt;00:00, 60.96it/s]"><pre class="notranslate">Selected optimization level O1:  Insert automatic casts around Pytorch functions and Tensor methods.

Defaults <span class="pl-k">for</span> this optimization level are:
enabled                <span class="pl-c1">:</span> True
opt_level              <span class="pl-c1">:</span> O1
cast_model_type        <span class="pl-c1">:</span> None
patch_torch_functions  <span class="pl-c1">:</span> True
keep_batchnorm_fp32    <span class="pl-c1">:</span> None
master_weights         <span class="pl-c1">:</span> None
loss_scale             <span class="pl-c1">:</span> dynamic
Processing user overrides (additional kwargs that are not None)...
After processing overrides, optimization options are:
enabled                <span class="pl-c1">:</span> True
opt_level              <span class="pl-c1">:</span> O1
cast_model_type        <span class="pl-c1">:</span> None
patch_torch_functions  <span class="pl-c1">:</span> True
keep_batchnorm_fp32    <span class="pl-c1">:</span> None
master_weights         <span class="pl-c1">:</span> None
loss_scale             <span class="pl-c1">:</span> dynamic
100%<span class="pl-k">|</span>██████████<span class="pl-k">|</span> 10/10 [00:<span class="pl-k">00&lt;</span>00:00, 60.96it/s]</pre></div>
<h2 dir="auto">Expected behavior</h2>

<p dir="auto">I hope these two features can be used at the same time.</p>
<h2 dir="auto">Environment</h2>
<div class="highlight highlight-source-shell notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="➜ python collect_env.py 
Collecting environment information...
PyTorch version: 1.6.0
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.4 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: Could not collect

Python version: 3.8 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 10.2.89
GPU models and configuration: 
GPU 0: GeForce RTX 2080 Ti
GPU 1: GeForce RTX 2080 Ti

Nvidia driver version: 440.100
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] efficientnet-pytorch==0.7.0
[pip3] numpy==1.19.1
[pip3] segmentation-models-pytorch==0.1.0
[pip3] torch==1.6.0
[pip3] torchvision==0.7.0
[conda] blas                      1.0                         mkl    defaults
[conda] cudatoolkit               10.2.89              hfd86e86_1    defaults
[conda] efficientnet-pytorch      0.7.0                    pypi_0    pypi
[conda] mkl                       2020.2                      256    defaults
[conda] mkl-service               2.3.0            py38he904b0f_0    defaults
[conda] mkl_fft                   1.1.0            py38h23d657b_0    defaults
[conda] mkl_random                1.1.1            py38h0573a6f_0    defaults
[conda] numpy                     1.19.1           py38hbc911f0_0    defaults
[conda] numpy-base                1.19.1           py38hfa32c7d_0    defaults
[conda] pytorch                   1.6.0           py3.8_cuda10.2.89_cudnn7.6.5_0    pytorch
[conda] segmentation-models-pytorch 0.1.0                    pypi_0    pypi
[conda] torchvision               0.7.0                py38_cu102    pytorch"><pre class="notranslate">➜ python collect_env.py 
Collecting environment information...
PyTorch version: 1.6.0
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.4 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: Could not collect

Python version: 3.8 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 10.2.89
GPU models and configuration: 
GPU 0: GeForce RTX 2080 Ti
GPU 1: GeForce RTX 2080 Ti

Nvidia driver version: 440.100
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] efficientnet-pytorch==0.7.0
[pip3] numpy==1.19.1
[pip3] segmentation-models-pytorch==0.1.0
[pip3] torch==1.6.0
[pip3] torchvision==0.7.0
[conda] blas                      1.0                         mkl    defaults
[conda] cudatoolkit               10.2.89              hfd86e86_1    defaults
[conda] efficientnet-pytorch      0.7.0                    pypi_0    pypi
[conda] mkl                       2020.2                      256    defaults
[conda] mkl-service               2.3.0            py38he904b0f_0    defaults
[conda] mkl_fft                   1.1.0            py38h23d657b_0    defaults
[conda] mkl_random                1.1.1            py38h0573a6f_0    defaults
[conda] numpy                     1.19.1           py38hbc911f0_0    defaults
[conda] numpy-base                1.19.1           py38hfa32c7d_0    defaults
[conda] pytorch                   1.6.0           py3.8_cuda10.2.89_cudnn7.6.5_0    pytorch
[conda] segmentation-models-pytorch 0.1.0                    pypi_0    pypi
[conda] torchvision               0.7.0                py38_cu102    pytorch</pre></div>
<h2 dir="auto">Additional context</h2>

<p dir="auto">None. Thanks.</p>
<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/mcarilli/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/mcarilli">@mcarilli</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">torch.utils.checkpoint.checkpoint fails when used with torch.cuda.amp on pytorch nightly. They both allow to reduce the memory usage so seem like a natural combination to use together.</p>
<h2 dir="auto">To Reproduce</h2>
<p dir="auto">Here is an example program derived from <a href="https://pytorch.org/docs/master/notes/amp_examples.html#typical-mixed-precision-training" rel="nofollow">https://pytorch.org/docs/master/notes/amp_examples.html#typical-mixed-precision-training</a></p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import torch
from torchvision.models import resnet18

amp_enabled = True
checkpoint_enabled = True
device = torch.device('cuda')

class ResNetWithCheckpoints(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.base = resnet18()

    def forward(self, x):
        base = self.base
        x = base.conv1(x)
        x = base.bn1(x)
        x = base.relu(x)
        x = base.maxpool(x)

        if checkpoint_enabled:
            x = torch.utils.checkpoint.checkpoint(base.layer1, x)
            x = torch.utils.checkpoint.checkpoint(base.layer2, x)
            x = torch.utils.checkpoint.checkpoint(base.layer3, x)
            x = torch.utils.checkpoint.checkpoint(base.layer4, x)
        else:
            x = base.layer1(x)
            x = base.layer2(x)
            x = base.layer3(x)
            x = base.layer4(x)

        x = base.avgpool(x)
        x = torch.flatten(x, 1)
        x = base.fc(x)

        return x

model = ResNetWithCheckpoints()
model.to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
loss_fn = torch.nn.CrossEntropyLoss()
scaler = torch.cuda.amp.GradScaler(enabled=amp_enabled)

xs = torch.randn(5, 3, 128, 128).to(device)
ys = torch.tensor(range(len(xs))).to(device)
for _ in range(20):
    optimizer.zero_grad()
    with torch.cuda.amp.autocast(enabled=amp_enabled):
        ys_pred = model(xs)
        loss = loss_fn(ys_pred, ys)
    print(float(loss))
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()"><pre class="notranslate"><code class="notranslate">import torch
from torchvision.models import resnet18

amp_enabled = True
checkpoint_enabled = True
device = torch.device('cuda')

class ResNetWithCheckpoints(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.base = resnet18()

    def forward(self, x):
        base = self.base
        x = base.conv1(x)
        x = base.bn1(x)
        x = base.relu(x)
        x = base.maxpool(x)

        if checkpoint_enabled:
            x = torch.utils.checkpoint.checkpoint(base.layer1, x)
            x = torch.utils.checkpoint.checkpoint(base.layer2, x)
            x = torch.utils.checkpoint.checkpoint(base.layer3, x)
            x = torch.utils.checkpoint.checkpoint(base.layer4, x)
        else:
            x = base.layer1(x)
            x = base.layer2(x)
            x = base.layer3(x)
            x = base.layer4(x)

        x = base.avgpool(x)
        x = torch.flatten(x, 1)
        x = base.fc(x)

        return x

model = ResNetWithCheckpoints()
model.to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
loss_fn = torch.nn.CrossEntropyLoss()
scaler = torch.cuda.amp.GradScaler(enabled=amp_enabled)

xs = torch.randn(5, 3, 128, 128).to(device)
ys = torch.tensor(range(len(xs))).to(device)
for _ in range(20):
    optimizer.zero_grad()
    with torch.cuda.amp.autocast(enabled=amp_enabled):
        ys_pred = model(xs)
        loss = loss_fn(ys_pred, ys)
    print(float(loss))
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
</code></pre></div>
<p dir="auto">This script fails when both amp and checkpointing are enabled, and does not fail otherwise:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="7.273253440856934
Traceback (most recent call last):
  File &quot;amp_checkpoint.py&quot;, line 51, in &lt;module&gt;
    scaler.scale(loss).backward()
  File &quot;lib/python3.6/site-packages/torch/tensor.py&quot;, line 184, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File &quot;lib/python3.6/site-packages/torch/autograd/__init__.py&quot;, line 115, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: Input type (torch.cuda.HalfTensor) and weight type (torch.cuda.FloatTensor) should be the same"><pre class="notranslate"><code class="notranslate">7.273253440856934
Traceback (most recent call last):
  File "amp_checkpoint.py", line 51, in &lt;module&gt;
    scaler.scale(loss).backward()
  File "lib/python3.6/site-packages/torch/tensor.py", line 184, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "lib/python3.6/site-packages/torch/autograd/__init__.py", line 115, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: Input type (torch.cuda.HalfTensor) and weight type (torch.cuda.FloatTensor) should be the same
</code></pre></div>
<h2 dir="auto">Expected behavior</h2>
<p dir="auto">I expect it to work with both amp and checkpoint enabled (I also expect memory savings from checkpointing but I'm not yet sure that my example should provide them).</p>
<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.6.0.dev20200502+cu101
Is debug build: No
CUDA used to build PyTorch: 10.1

OS: Ubuntu 18.04.4 LTS
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
CMake version: Could not collect

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: 10.1.243
GPU models and configuration: 
GPU 0: GeForce GTX 1050 Ti
GPU 1: GeForce RTX 2080 Ti   (this is the one used here)

Nvidia driver version: 430.50
cuDNN version: /usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn.so.7

Versions of relevant libraries:
[pip3] numpy==1.18.3
[pip3] torch==1.6.0.dev20200502+cu101
[pip3] torchvision==0.7.0.dev20200501+cu101
[conda] Could not collect"><pre class="notranslate"><code class="notranslate">Collecting environment information...
PyTorch version: 1.6.0.dev20200502+cu101
Is debug build: No
CUDA used to build PyTorch: 10.1

OS: Ubuntu 18.04.4 LTS
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
CMake version: Could not collect

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: 10.1.243
GPU models and configuration: 
GPU 0: GeForce GTX 1050 Ti
GPU 1: GeForce RTX 2080 Ti   (this is the one used here)

Nvidia driver version: 430.50
cuDNN version: /usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn.so.7

Versions of relevant libraries:
[pip3] numpy==1.18.3
[pip3] torch==1.6.0.dev20200502+cu101
[pip3] torchvision==0.7.0.dev20200501+cu101
[conda] Could not collect
</code></pre></div>