<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">This bug report is created as instructed by the fail message below:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="Traceback (most recent call last):
File &quot;qat_residual.py&quot;, line 60, in &lt;module&gt;
TRAINING_DEVICE, print_freq=500)
File &quot;/home/lijun/repo/quant/utils.py&quot;, line 73, in train_one_epoch
output = model(image)
File &quot;/home/lijun/conda/lib/python3.7/site-packages/torch/nn/modules/module.py&quot;, line 550, in __call__
result = self.forward(*input, **kwargs)
File &quot;/home/lijun/repo/quant/residual_model.py&quot;, line 312, in forward
return self._forward_impl(x)
File &quot;/home/lijun/repo/quant/residual_model.py&quot;, line 295, in _forward_impl
x = self.conv1(x)
File &quot;/home/lijun/conda/lib/python3.7/site-packages/torch/nn/modules/module.py&quot;, line 550, in __call__
result = self.forward(*input, **kwargs)
File &quot;/home/lijun/conda/lib/python3.7/site-packages/torch/nn/qat/modules/conv.py&quot;, line 36, in forward
self._conv_forward(input, self.weight_fake_quant(self.weight)))
File &quot;/home/lijun/conda/lib/python3.7/site-packages/torch/nn/modules/module.py&quot;, line 550, in __call__
result = self.forward(*input, **kwargs)
File &quot;/home/lijun/conda/lib/python3.7/site-packages/torch/quantization/fake_quantize.py&quot;, line 86, in forward
self.ch_axis, self.quant_min, self.quant_max)
RuntimeError: iter.device(arg).is_cuda() INTERNAL ASSERT FAILED at /opt/conda/conda-bld/pytorch_1587428398394/work/aten/src/ATen/native/cuda/Loops.cuh:56, please report a bug to PyTorch."><pre class="notranslate"><code class="notranslate">Traceback (most recent call last):
File "qat_residual.py", line 60, in &lt;module&gt;
TRAINING_DEVICE, print_freq=500)
File "/home/lijun/repo/quant/utils.py", line 73, in train_one_epoch
output = model(image)
File "/home/lijun/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
result = self.forward(*input, **kwargs)
File "/home/lijun/repo/quant/residual_model.py", line 312, in forward
return self._forward_impl(x)
File "/home/lijun/repo/quant/residual_model.py", line 295, in _forward_impl
x = self.conv1(x)
File "/home/lijun/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
result = self.forward(*input, **kwargs)
File "/home/lijun/conda/lib/python3.7/site-packages/torch/nn/qat/modules/conv.py", line 36, in forward
self._conv_forward(input, self.weight_fake_quant(self.weight)))
File "/home/lijun/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
result = self.forward(*input, **kwargs)
File "/home/lijun/conda/lib/python3.7/site-packages/torch/quantization/fake_quantize.py", line 86, in forward
self.ch_axis, self.quant_min, self.quant_max)
RuntimeError: iter.device(arg).is_cuda() INTERNAL ASSERT FAILED at /opt/conda/conda-bld/pytorch_1587428398394/work/aten/src/ATen/native/cuda/Loops.cuh:56, please report a bug to PyTorch.
</code></pre></div>
<h2 dir="auto">To Reproduce</h2>
<p dir="auto">Steps to reproduce the behavior:</p>
<ol dir="auto">
<li>A float-point number ResNet was trained</li>
<li>The pre-trained ResNet was prepared for QAT</li>
<li>When QAT ran on the first epoch, it raised the message above.</li>
</ol>

<h2 dir="auto">Expected behavior</h2>

<p dir="auto">The model should be able to move from system memory to GPU memory.</p>
<h2 dir="auto">Environment</h2>
<p dir="auto">Please copy and paste the output from our<br>
<a href="https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py" rel="nofollow">environment collection script</a><br>
(or fill out the checklist below manually).</p>
<p dir="auto">You can get the script and run it with:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="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"><pre class="notranslate"><code class="notranslate">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
</code></pre></div>
<ul dir="auto">
<li>PyTorch Version (e.g., 1.0):</li>
<li>OS (e.g., Linux):</li>
<li>How you installed PyTorch (<code class="notranslate">conda</code>, <code class="notranslate">pip</code>, source):</li>
<li>Build command you used (if compiling from source):</li>
<li>Python version:</li>
<li>CUDA/cuDNN version:</li>
<li>GPU models and configuration:</li>
<li>Any other relevant information:</li>
</ul>
<h2 dir="auto">Additional context</h2>

<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/jerryzh168/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/jerryzh168">@jerryzh168</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/jianyuh/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/jianyuh">@jianyuh</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/dzhulgakov/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/dzhulgakov">@dzhulgakov</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/raghuramank100/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/raghuramank100">@raghuramank100</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/jamesr66a/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/jamesr66a">@jamesr66a</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">I'm following this tutorial : <a href="https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html" rel="nofollow">https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html</a><br>
This tutorial uses a CPU, it works fine with a CPU, but to speed it up I want to move it to the GPU and I get an error. It works with Pytorch version 1.4 but not with Pytorch nightly (1.5.0.dev20200413).</p>
<h2 dir="auto">To Reproduce</h2>
<p dir="auto">Steps to reproduce the behavior:</p>
<p dir="auto">1.Follow the following tutorial:<br>
<a href="https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html" rel="nofollow">https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html</a></p>
<ol start="2" dir="auto">
<li>In the "5. Quantization-aware training" section</li>
</ol>
<p dir="auto">Instead of doing</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="for nepoch in range(8):
    train_one_epoch(qat_model, criterion, optimizer, data_loader, torch.device('cpu'), num_train_batches)"><pre class="notranslate"><code class="notranslate">for nepoch in range(8):
    train_one_epoch(qat_model, criterion, optimizer, data_loader, torch.device('cpu'), num_train_batches)
</code></pre></div>
<p dir="auto">do :</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="device = torch.device(&quot;cuda:0&quot;)
for nepoch in range(8):
    device = torch.device(&quot;cuda:0&quot;)
    train_one_epoch(qat_model, criterion, optimizer, train_loader, torch.device(device), num_train_batches)"><pre class="notranslate"><code class="notranslate">device = torch.device("cuda:0")
for nepoch in range(8):
    device = torch.device("cuda:0")
    train_one_epoch(qat_model, criterion, optimizer, train_loader, torch.device(device), num_train_batches)
</code></pre></div>
<p dir="auto">also changed in the section "2. Helper functions":</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="def load_model(model_file):
    model = MobileNetV2()
    state_dict = torch.load(model_file)
    model.load_state_dict(state_dict)
    model.to(device) #device instead of 'cpu'
    return model"><pre class="notranslate"><code class="notranslate">def load_model(model_file):
    model = MobileNetV2()
    state_dict = torch.load(model_file)
    model.load_state_dict(state_dict)
    model.to(device) #device instead of 'cpu'
    return model
</code></pre></div>
<p dir="auto">3.Get the error</p>
<div class="highlight highlight-text-shell-session notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="F:\Anaconda3\envs\pytorch-nightly\lib\site-packages\torch\quantization\fake_quantize.py in forward(self, X)
     84             if self.qscheme == torch.per_channel_symmetric or self.qscheme == torch.per_channel_affine:
     85                 X = torch.fake_quantize_per_channel_affine(X, self.scale, self.zero_point,
---&gt; 86                                                            self.ch_axis, self.quant_min, self.quant_max)
     87             else:
     88                 X = torch.fake_quantize_per_tensor_affine(X, float(self.scale),

RuntimeError: iter.device(arg).is_cuda() INTERNAL ASSERT FAILED at C:/cb/pytorch_1000000000000/work/aten/src\ATen/native/cuda/Loops.cuh:60, please report a bug to PyTorch. "><pre class="notranslate"><span class="pl-c1">F:\Anaconda3\envs\pytorch-nightly\lib\site-packages\torch\quantization\fake_quantize.py in forward(self, X)</span>
<span class="pl-c1">     84             if self.qscheme == torch.per_channel_symmetric or self.qscheme == torch.per_channel_affine:</span>
<span class="pl-c1">     85                 X = torch.fake_quantize_per_channel_affine(X, self.scale, self.zero_point,</span>
<span class="pl-c1">---&gt; 86                                                            self.ch_axis, self.quant_min, self.quant_max)</span>
<span class="pl-c1">     87             else:</span>
<span class="pl-c1">     88                 X = torch.fake_quantize_per_tensor_affine(X, float(self.scale),</span>

<span class="pl-c1">RuntimeError: iter.device(arg).is_cuda() INTERNAL ASSERT FAILED at C:/cb/pytorch_1000000000000/work/aten/src\ATen/native/cuda/Loops.cuh:60, please report a bug to PyTorch. </span></pre></div>
<h2 dir="auto">Expected behavior</h2>
<p dir="auto">No error, run smoothly. QAT is supposed to support CUDA.<br>
Quote from the documentation<br>
"Quantization-aware training (through FakeQuantize) supports both CPU and CUDA."</p>
<p dir="auto">However, it works with  PyTorch Version : 1.4.0.</p>
<h2 dir="auto">Environment</h2>
<ul dir="auto">
<li>PyTorch Version : 1.5.0.dev20200413 (pytorch nightly)</li>
<li>OS :Windows 10</li>
<li>How you installed PyTorch :conda (version 4.7.12)</li>
<li>Build command you used (if compiling from source):<br>
conda install pytorch torchvision cudatoolkit=10.1 -c pytorch-nightly -c defaults -c conda-forge</li>
<li>Python version: 3.7.7</li>
<li>CUDA/cuDNN version: 10.1</li>
<li>GPU models and configuration: GPU: Nvidia Geforce RTX 2060 Super (8Go GDDR6)</li>
</ul>
<p dir="auto">Versions of relevant libraries:<br>
[pip] numpy==1.18.1<br>
[pip] torch==1.5.0.dev20200413<br>
[pip] torchvision==0.7.0.dev20200415<br>
[conda] blas                      1.0                         mkl<br>
[conda] cudatoolkit               10.1.243             h74a9793_0<br>
[conda] mkl                       2020.0                      166<br>
[conda] mkl-service               2.3.0            py37hb782905_0<br>
[conda] mkl_fft                   1.0.15           py37h14836fe_0<br>
[conda] mkl_random                1.1.0            py37h675688f_0<br>
[conda] numpy                     1.18.1           py37h93ca92e_0<br>
[conda] numpy-base                1.18.1           py37hc3f5095_1<br>
[conda] pytorch                   1.5.0.dev20200413 py3.7_cuda101_cudnn7_0    pytorch-nightly<br>
[conda] torchvision               0.7.0.dev20200415      py37_cu101    pytorch-nightly</p>
<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/jerryzh168/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/jerryzh168">@jerryzh168</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/jianyuh/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/jianyuh">@jianyuh</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/dzhulgakov/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/dzhulgakov">@dzhulgakov</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/raghuramank100/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/raghuramank100">@raghuramank100</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/jamesr66a/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/jamesr66a">@jamesr66a</a></p>