Issue 1:
I have a DenseNet-like model for speech recognition (https://arxiv.org/abs/1904.03288). My code: https://github.com/vadimkantorov/convasr/blob/master/models.py#L94

When I export it to ONNX as is, everything works okay. When I apply model = apex.amp.initialize(model, opt_level = 'O3') prior to export (basically to have model params converted to float16 and have input/output casts from float32), I get this:

RuntimeError: Cannot insert a Tensor that requires grad as a constant. Consider making it a parameter or input, or detaching the gradient

The error is really strange.

When I cast model parameters to float16 manually and have input cast to float16 manually, I have this bug instead: microsoft/onnxruntime#4405

Full stack trace:

Traceback (most recent call last):
File "train.py", line 383, in <module>
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

cc @mcarilli @houseroad @spandantiwari @lara-hdr @BowenBao @neginraoof

Issue 2:
🐛 Bug

Traceback (most recent call last):
File "repro_autocast_trace.py", line 33, in <module>
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

To Reproduce

Run the following script:

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))

The issue looks similar to #40910, which was using tracing with apex amp.

cc @mcarilli @ptrblck @gmagogsfm @mruberry, @ngimel