Cannot export O2 and O3 model with ONNX

Hello, I recently started using Apex mixed precision training. My first goal is to reduce my final model size on disk. For that purpose, I train my model with the O2 optimization level. I then require to export this model using PyTorch ONNX integration, but it fails with the following error:

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

My network is a simple ResNet, and it fails at the very first layer (a Conv2D). Note that it doesn't fail with O1 optimization level, and that the training also doesn't fail, it really just is exporting that doesn't work.

Here is the full trace if needed:

Traceback (most recent call last):
  File "[...]\Anaconda3\envs\PyTorch 1.4\lib\site-packages\IPython\core\interactiveshell.py", line 2722, in safe_execfile
    self.compile if shell_futures else None)
  File "[...]\Anaconda3\envs\PyTorch 1.4\lib\site-packages\IPython\utils\py3compat.py", line 168, in execfile
    exec(compiler(f.read(), fname, 'exec'), glob, loc)
  File "[...]\train.py", line 167, in <module>
    trainer.train(model)
  File "[...]\model.py", line 644, in train
    model.save(self.output_path, prefix=("ep%i_%s" % (epoch, self.nndef_name)))
  File "[...]\model.py", line 144, in save
    torch.onnx.export(self.encoder, dummy_input, path + "/" + prefix + "_NexE.onnx", verbose=False)
  File "[...]\Anaconda3\envs\PyTorch 1.4\lib\site-packages\torch\onnx\__init__.py", line 148, in export
    strip_doc_string, dynamic_axes, keep_initializers_as_inputs)
  File "[...]\Anaconda3\envs\PyTorch 1.4\lib\site-packages\torch\onnx\utils.py", line 66, in export
    dynamic_axes=dynamic_axes, keep_initializers_as_inputs=keep_initializers_as_inputs)
  File "[...]\Anaconda3\envs\PyTorch 1.4\lib\site-packages\torch\onnx\utils.py", line 416, in _export
    fixed_batch_size=fixed_batch_size)
  File "[...]\Anaconda3\envs\PyTorch 1.4\lib\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 "[...]\Anaconda3\envs\PyTorch 1.4\lib\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 "[...]\Anaconda3\envs\PyTorch 1.4\lib\site-packages\torch\jit\__init__.py", line 277, in _get_trace_graph
    outs = ONNXTracedModule(f, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
  File "[...]\Anaconda3\envs\PyTorch 1.4\lib\site-packages\torch\nn\modules\module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "[...]\Anaconda3\envs\PyTorch 1.4\lib\site-packages\torch\jit\__init__.py", line 360, in forward
    self._force_outplace,
  File "[...]\Anaconda3\envs\PyTorch 1.4\lib\site-packages\torch\jit\__init__.py", line 347, in wrapper
    outs.append(self.inner(*trace_inputs))
  File "[...]\Anaconda3\envs\PyTorch 1.4\lib\site-packages\torch\nn\modules\module.py", line 530, in __call__
    result = self._slow_forward(*input, **kwargs)
  File "[...]\Anaconda3\envs\PyTorch 1.4\lib\site-packages\torch\nn\modules\module.py", line 516, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "[...]\Anaconda3\envs\PyTorch 1.4\lib\site-packages\apex\amp\_initialize.py", line 197, in new_fwd
    **applier(kwargs, input_caster))
  File "C:\WORK\GIT\com.unity.smartassets\MachineLearning\.training\layout_generator\v1_0_0\models\MobileResNet.py", line 153, in forward
    x = self.conv1(x)
  File "[...]\Anaconda3\envs\PyTorch 1.4\lib\site-packages\torch\nn\modules\module.py", line 530, in __call__
    result = self._slow_forward(*input, **kwargs)
  File "[...]\Anaconda3\envs\PyTorch 1.4\lib\site-packages\torch\nn\modules\module.py", line 516, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "[...]\Anaconda3\envs\PyTorch 1.4\lib\site-packages\torch\nn\modules\conv.py", line 345, in forward
    return self.conv2d_forward(input, self.weight)
  File "[...]\Anaconda3\envs\PyTorch 1.4\lib\site-packages\torch\nn\modules\conv.py", line 342, in conv2d_forward
    self.padding, self.dilation, self.groups)

I'm using PyTorch 1.4 with Cuda 10.1

Start with apex/amp/_initialize.py, where the traceback enters the AMP-wrapped forward, and the torch.onnx.export call described in the report. Reproduce on PyTorch 1.4 with the simple ResNet using O2, then compare with O1. Done means the O2 model exports successfully to ONNX without the reported gradient-constant error.