<p dir="auto">I have a DenseNet-like model for speech recognition (<a href="https://arxiv.org/abs/1904.03288" rel="nofollow">https://arxiv.org/abs/1904.03288</a>). My code:  <a href="https://github.com/vadimkantorov/convasr/blob/master/models.py#L94">https://github.com/vadimkantorov/convasr/blob/master/models.py#L94</a></p>
<p dir="auto">When I export it to ONNX as is, everything works okay. When I apply <code class="notranslate">model = apex.amp.initialize(model, opt_level = 'O3')</code> prior to export (basically to have model params converted to float16 and have input/output casts from float32), I get this:</p>
<p dir="auto"><code class="notranslate">RuntimeError: Cannot insert a Tensor that requires grad as a constant. Consider making it a parameter or input, or detaching the gradient</code></p>
<p dir="auto">The error is really strange.</p>
<p dir="auto">When I cast model parameters to float16 manually and have input cast to float16 manually, I have this bug instead: <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="649880059" data-permission-text="Title is private" data-url="https://github.com/microsoft/onnxruntime/issues/4405" data-hovercard-type="issue" data-hovercard-url="/microsoft/onnxruntime/issues/4405/hovercard" href="https://github.com/microsoft/onnxruntime/issues/4405">microsoft/onnxruntime#4405</a></p>
<p dir="auto">Full stack trace:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="Traceback (most recent call last):
  File &quot;train.py&quot;, line 383, in &lt;module&gt;
    main(parser.parse_args())
  File &quot;train.py&quot;, line 72, in main
    torch.onnx.export(model, (waveform_input,), args.onnx, opset_version = args.onnx_opset, export_params = args.onnx_export_params, do_constant_folding = True, input_names = ['x'], output_names = ['logits'], dynamic_axes = dict(x = {0 : 'B', 1 : 'T'}, logits = {0 : 'B', 2: 't'}))
  File &quot;/miniconda/lib/python3.7/site-packages/torch/onnx/__init__.py&quot;, line 148, in export
    strip_doc_string, dynamic_axes, keep_initializers_as_inputs)
  File &quot;/miniconda/lib/python3.7/site-packages/torch/onnx/utils.py&quot;, line 66, in export
    dynamic_axes=dynamic_axes, keep_initializers_as_inputs=keep_initializers_as_inputs)
  File &quot;/miniconda/lib/python3.7/site-packages/torch/onnx/utils.py&quot;, line 416, in _export
    fixed_batch_size=fixed_batch_size)
  File &quot;/miniconda/lib/python3.7/site-packages/torch/onnx/utils.py&quot;, line 279, in _model_to_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args, training)
  File &quot;/miniconda/lib/python3.7/site-packages/torch/onnx/utils.py&quot;, line 236, in _trace_and_get_graph_from_model
    trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph(model, args, _force_outplace=True, _return_inputs_states=True)
  File &quot;/miniconda/lib/python3.7/site-packages/torch/jit/__init__.py&quot;, line 272, in _get_trace_graph
    outs = ONNXTracedModule(f, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
  File &quot;/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py&quot;, line 532, in __call__
    result = self.forward(*input, **kwargs)
  File &quot;/miniconda/lib/python3.7/site-packages/torch/jit/__init__.py&quot;, line 355, in forward
    self._force_outplace,
  File &quot;/miniconda/lib/python3.7/site-packages/torch/jit/__init__.py&quot;, line 342, in wrapper
    outs.append(self.inner(*trace_inputs))
  File &quot;/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py&quot;, line 530, in __call__
    result = self._slow_forward(*input, **kwargs)
  File &quot;/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py&quot;, line 516, in _slow_forward
    result = self.forward(*input, **kwargs)
  File &quot;/miniconda/lib/python3.7/site-packages/apex/amp/_initialize.py&quot;, line 197, in new_fwd
    **applier(kwargs, input_caster))
  File &quot;/root/convasr/models.py&quot;, line 146, in forward
    x = subblock(x, residual = residual, lengths_fraction = xlen)
  File &quot;/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py&quot;, line 530, in __call__
    result = self._slow_forward(*input, **kwargs)
  File &quot;/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py&quot;, line 516, in _slow_forward
    result = self.forward(*input, **kwargs)
  File &quot;/root/convasr/models.py&quot;, line 74, in forward
    x = self.activation(bn(conv(x)), residual = residual_inputs)
  File &quot;/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py&quot;, line 530, in __call__
    result = self._slow_forward(*input, **kwargs)
  File &quot;/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py&quot;, line 516, in _slow_forward
    result = self.forward(*input, **kwargs)
  File &quot;/miniconda/lib/python3.7/site-packages/torch/nn/modules/container.py&quot;, line 100, in forward
    input = module(input)
  File &quot;/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py&quot;, line 530, in __call__
    result = self._slow_forward(*input, **kwargs)
  File &quot;/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py&quot;, line 516, in _slow_forward
    result = self.forward(*input, **kwargs)
  File &quot;/miniconda/lib/python3.7/site-packages/torch/nn/modules/conv.py&quot;, line 202, in forward
    self.padding, self.dilation, self.groups)
RuntimeError: Cannot insert a Tensor that requires grad as a constant. Consider making it a parameter or input, or detaching the gradient"><pre class="notranslate"><code class="notranslate">Traceback (most recent call last):
  File "train.py", line 383, in &lt;module&gt;
    main(parser.parse_args())
  File "train.py", line 72, in main
    torch.onnx.export(model, (waveform_input,), args.onnx, opset_version = args.onnx_opset, export_params = args.onnx_export_params, do_constant_folding = True, input_names = ['x'], output_names = ['logits'], dynamic_axes = dict(x = {0 : 'B', 1 : 'T'}, logits = {0 : 'B', 2: 't'}))
  File "/miniconda/lib/python3.7/site-packages/torch/onnx/__init__.py", line 148, in export
    strip_doc_string, dynamic_axes, keep_initializers_as_inputs)
  File "/miniconda/lib/python3.7/site-packages/torch/onnx/utils.py", line 66, in export
    dynamic_axes=dynamic_axes, keep_initializers_as_inputs=keep_initializers_as_inputs)
  File "/miniconda/lib/python3.7/site-packages/torch/onnx/utils.py", line 416, in _export
    fixed_batch_size=fixed_batch_size)
  File "/miniconda/lib/python3.7/site-packages/torch/onnx/utils.py", line 279, in _model_to_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args, training)
  File "/miniconda/lib/python3.7/site-packages/torch/onnx/utils.py", line 236, in _trace_and_get_graph_from_model
    trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph(model, args, _force_outplace=True, _return_inputs_states=True)
  File "/miniconda/lib/python3.7/site-packages/torch/jit/__init__.py", line 272, in _get_trace_graph
    outs = ONNXTracedModule(f, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
  File "/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/miniconda/lib/python3.7/site-packages/torch/jit/__init__.py", line 355, in forward
    self._force_outplace,
  File "/miniconda/lib/python3.7/site-packages/torch/jit/__init__.py", line 342, in wrapper
    outs.append(self.inner(*trace_inputs))
  File "/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 530, in __call__
    result = self._slow_forward(*input, **kwargs)
  File "/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 516, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/miniconda/lib/python3.7/site-packages/apex/amp/_initialize.py", line 197, in new_fwd
    **applier(kwargs, input_caster))
  File "/root/convasr/models.py", line 146, in forward
    x = subblock(x, residual = residual, lengths_fraction = xlen)
  File "/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 530, in __call__
    result = self._slow_forward(*input, **kwargs)
  File "/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 516, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/root/convasr/models.py", line 74, in forward
    x = self.activation(bn(conv(x)), residual = residual_inputs)
  File "/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 530, in __call__
    result = self._slow_forward(*input, **kwargs)
  File "/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 516, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/miniconda/lib/python3.7/site-packages/torch/nn/modules/container.py", line 100, in forward
    input = module(input)
  File "/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 530, in __call__
    result = self._slow_forward(*input, **kwargs)
  File "/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 516, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/miniconda/lib/python3.7/site-packages/torch/nn/modules/conv.py", line 202, in forward
    self.padding, self.dilation, self.groups)
RuntimeError: Cannot insert a Tensor that requires grad as a constant. Consider making it a parameter or input, or detaching the gradient
</code></pre></div>
<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> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/houseroad/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/houseroad">@houseroad</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/spandantiwari/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/spandantiwari">@spandantiwari</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/lara-hdr/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/lara-hdr">@lara-hdr</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/BowenBao/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/BowenBao">@BowenBao</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/neginraoof/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/neginraoof">@neginraoof</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>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="Traceback (most recent call last):
  File &quot;repro_autocast_trace.py&quot;, line 33, in &lt;module&gt;
    model = torch.jit.trace(model, (x, y))
  File &quot;/home/bowbao/repos/pytorch/torch/jit/_trace.py&quot;, line 742, in trace
    _module_class,
  File &quot;/home/bowbao/repos/pytorch/torch/jit/_trace.py&quot;, line 966, in trace_module
    _module_class,
  File &quot;/home/bowbao/repos/pytorch/torch/autograd/grad_mode.py&quot;, line 27, in decorate_context
    return func(*args, **kwargs)
  File &quot;/home/bowbao/repos/pytorch/torch/jit/_trace.py&quot;, line 335, in _check_trace
    _compilation_unit=torch._C.CompilationUnit(),
  File &quot;/home/bowbao/repos/pytorch/torch/jit/_trace.py&quot;, line 940, in trace_module
    _force_outplace,
  File &quot;/home/bowbao/repos/pytorch/torch/nn/modules/module.py&quot;, line 878, in _call_impl
    result = self._slow_forward(*input, **kwargs)
  File &quot;/home/bowbao/repos/pytorch/torch/nn/modules/module.py&quot;, line 851, in _slow_forward
    result = self.forward(*input, **kwargs)
  File &quot;repro_autocast_trace.py&quot;, line 14, in forward
    out = self.fc1(x)
  File &quot;/home/bowbao/repos/pytorch/torch/nn/modules/module.py&quot;, line 878, in _call_impl
    result = self._slow_forward(*input, **kwargs)
  File &quot;/home/bowbao/repos/pytorch/torch/nn/modules/module.py&quot;, line 851, in _slow_forward
    result = self.forward(*input, **kwargs)
  File &quot;/home/bowbao/repos/pytorch/torch/nn/modules/linear.py&quot;, line 94, in forward
    return F.linear(input, self.weight, self.bias)
  File &quot;/home/bowbao/repos/pytorch/torch/nn/functional.py&quot;, line 1665, in linear
    ret = torch.addmm(bias, input, weight.t())
RuntimeError: Cannot insert a Tensor that requires grad as a constant. Consider making it a parameter or input, or detaching the gradient"><pre class="notranslate"><code class="notranslate">Traceback (most recent call last):
  File "repro_autocast_trace.py", line 33, in &lt;module&gt;
    model = torch.jit.trace(model, (x, y))
  File "/home/bowbao/repos/pytorch/torch/jit/_trace.py", line 742, in trace
    _module_class,
  File "/home/bowbao/repos/pytorch/torch/jit/_trace.py", line 966, in trace_module
    _module_class,
  File "/home/bowbao/repos/pytorch/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/bowbao/repos/pytorch/torch/jit/_trace.py", line 335, in _check_trace
    _compilation_unit=torch._C.CompilationUnit(),
  File "/home/bowbao/repos/pytorch/torch/jit/_trace.py", line 940, in trace_module
    _force_outplace,
  File "/home/bowbao/repos/pytorch/torch/nn/modules/module.py", line 878, in _call_impl
    result = self._slow_forward(*input, **kwargs)
  File "/home/bowbao/repos/pytorch/torch/nn/modules/module.py", line 851, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "repro_autocast_trace.py", line 14, in forward
    out = self.fc1(x)
  File "/home/bowbao/repos/pytorch/torch/nn/modules/module.py", line 878, in _call_impl
    result = self._slow_forward(*input, **kwargs)
  File "/home/bowbao/repos/pytorch/torch/nn/modules/module.py", line 851, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/home/bowbao/repos/pytorch/torch/nn/modules/linear.py", line 94, in forward
    return F.linear(input, self.weight, self.bias)
  File "/home/bowbao/repos/pytorch/torch/nn/functional.py", line 1665, in linear
    ret = torch.addmm(bias, input, weight.t())
RuntimeError: Cannot insert a Tensor that requires grad as a constant. Consider making it a parameter or input, or detaching the gradient
</code></pre></div>
<h2 dir="auto">To Reproduce</h2>
<p dir="auto">Run the following script:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import torch
import torch.nn as nn
import torch.nn.functional as F

class NeuralNetWithLoss(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(NeuralNetWithLoss, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, num_classes)

    def forward(self, x, target):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return F.nll_loss(F.log_softmax(out, dim=1), target), out

input_size = 784
hidden_size = 500
num_classes = 10

model = NeuralNetWithLoss(input_size, hidden_size, num_classes)
model.cuda()
model.eval()

with torch.cuda.amp.autocast():
    x = torch.randn(32, input_size).cuda()
    y = torch.randint(0, 10, (32,)).cuda()

    model = torch.jit.trace(model, (x, y))"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">torch</span>
<span class="pl-k">import</span> <span class="pl-s1">torch</span>.<span class="pl-s1">nn</span> <span class="pl-k">as</span> <span class="pl-s1">nn</span>
<span class="pl-k">import</span> <span class="pl-s1">torch</span>.<span class="pl-s1">nn</span>.<span class="pl-s1">functional</span> <span class="pl-k">as</span> <span class="pl-v">F</span>

<span class="pl-k">class</span> <span class="pl-v">NeuralNetWithLoss</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-s1">input_size</span>, <span class="pl-s1">hidden_size</span>, <span class="pl-s1">num_classes</span>):
        <span class="pl-en">super</span>(<span class="pl-v">NeuralNetWithLoss</span>, <span class="pl-s1">self</span>).<span class="pl-en">__init__</span>()
        <span class="pl-s1">self</span>.<span class="pl-s1">fc1</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">Linear</span>(<span class="pl-s1">input_size</span>, <span class="pl-s1">hidden_size</span>)
        <span class="pl-s1">self</span>.<span class="pl-s1">relu</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">ReLU</span>()
        <span class="pl-s1">self</span>.<span class="pl-s1">fc2</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">Linear</span>(<span class="pl-s1">hidden_size</span>, <span class="pl-s1">num_classes</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">target</span>):
        <span class="pl-s1">out</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">fc1</span>(<span class="pl-s1">x</span>)
        <span class="pl-s1">out</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">relu</span>(<span class="pl-s1">out</span>)
        <span class="pl-s1">out</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">fc2</span>(<span class="pl-s1">out</span>)
        <span class="pl-k">return</span> <span class="pl-v">F</span>.<span class="pl-en">nll_loss</span>(<span class="pl-v">F</span>.<span class="pl-en">log_softmax</span>(<span class="pl-s1">out</span>, <span class="pl-s1">dim</span><span class="pl-c1">=</span><span class="pl-c1">1</span>), <span class="pl-s1">target</span>), <span class="pl-s1">out</span>

<span class="pl-s1">input_size</span> <span class="pl-c1">=</span> <span class="pl-c1">784</span>
<span class="pl-s1">hidden_size</span> <span class="pl-c1">=</span> <span class="pl-c1">500</span>
<span class="pl-s1">num_classes</span> <span class="pl-c1">=</span> <span class="pl-c1">10</span>

<span class="pl-s1">model</span> <span class="pl-c1">=</span> <span class="pl-v">NeuralNetWithLoss</span>(<span class="pl-s1">input_size</span>, <span class="pl-s1">hidden_size</span>, <span class="pl-s1">num_classes</span>)
<span class="pl-s1">model</span>.<span class="pl-en">cuda</span>()
<span class="pl-s1">model</span>.<span class="pl-en">eval</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-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">randn</span>(<span class="pl-c1">32</span>, <span class="pl-s1">input_size</span>).<span class="pl-en">cuda</span>()
    <span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">randint</span>(<span class="pl-c1">0</span>, <span class="pl-c1">10</span>, (<span class="pl-c1">32</span>,)).<span class="pl-en">cuda</span>()

    <span class="pl-s1">model</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-s1">jit</span>.<span class="pl-en">trace</span>(<span class="pl-s1">model</span>, (<span class="pl-s1">x</span>, <span class="pl-s1">y</span>))</pre></div>
<p dir="auto">The issue looks similar to <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="649861898" data-permission-text="Title is private" data-url="https://github.com/pytorch/pytorch/issues/40910" data-hovercard-type="issue" data-hovercard-url="/pytorch/pytorch/issues/40910/hovercard" href="https://github.com/pytorch/pytorch/issues/40910">#40910</a>, which was using tracing with apex amp.</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> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/ptrblck/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/ptrblck">@ptrblck</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/gmagogsfm/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/gmagogsfm">@gmagogsfm</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/mruberry/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/mruberry">@mruberry</a>, <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/ngimel/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/ngimel">@ngimel</a></p>