"Cannot insert a Tensor that requires grad as a constant" error when exporting an onnx model

I have a model that I successfully export using torch.onnx.export. When I introduce AMP into the code, export throws an exception that says "Cannot insert a Tensor that requires grad as a constant".

Here is some code that triggers the issue:

import torch
from apex import amp

class Foo(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Conv2d(1, 1, kernel_size=5, stride=1, padding=2)

    def forward(self, x):
        return self.layer(x)

model = Foo().cuda()
optimizer = torch.optim.AdamW(model.parameters())
# Removing the following line fixes the issue.
model, optimizer = amp.initialize(model, optimizer, opt_level="O1")

dummy_input = torch.randn(1, 1, 28, 28).cuda()
# Alternatively, removing the following line fixes the issue.
output = model(dummy_input)
torch.onnx.export(model, dummy_input, './amp.onnx')

As I mention in the comments above, removing the line that makes the forward pass before exporting fixes the issue. Alternatively, removing the call to AMP also fixes the issue.

The full error message is:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-4-efcde020fc8f> in <module>
      1 dummy_input = torch.randn(1, 1, 28, 28).cuda()
      2 output = model(dummy_input) # alternatively, removing this line fixes the issue
----> 3 torch.onnx.export(model, dummy_input, './amp.onnx')

~/anaconda3/envs/amp/lib/python3.7/site-packages/torch/onnx/__init__.py in export(model, args, f, export_params, verbose, training, input_names, output_names, aten, export_raw_ir, operator_export_type, opset_version, _retain_param_name, do_constant_folding, example_outputs, strip_doc_string, dynamic_axes, keep_initializers_as_inputs)
    146                         operator_export_type, opset_version, _retain_param_name,
    147                         do_constant_folding, example_outputs,
--> 148                         strip_doc_string, dynamic_axes, keep_initializers_as_inputs)
    149 
    150 

~/anaconda3/envs/amp/lib/python3.7/site-packages/torch/onnx/utils.py in export(model, args, f, export_params, verbose, training, input_names, output_names, aten, export_raw_ir, operator_export_type, opset_version, _retain_param_name, do_constant_folding, example_outputs, strip_doc_string, dynamic_axes, keep_initializers_as_inputs)
     64             _retain_param_name=_retain_param_name, do_constant_folding=do_constant_folding,
     65             example_outputs=example_outputs, strip_doc_string=strip_doc_string,
---> 66             dynamic_axes=dynamic_axes, keep_initializers_as_inputs=keep_initializers_as_inputs)
     67 
     68 

~/anaconda3/envs/amp/lib/python3.7/site-packages/torch/onnx/utils.py in _export(model, args, f, export_params, verbose, training, input_names, output_names, operator_export_type, export_type, example_outputs, propagate, opset_version, _retain_param_name, do_constant_folding, strip_doc_string, dynamic_axes, keep_initializers_as_inputs, fixed_batch_size)
    414                                                         example_outputs, propagate,
    415                                                         _retain_param_name, do_constant_folding,
--> 416                                                         fixed_batch_size=fixed_batch_size)
    417 
    418         # TODO: Don't allocate a in-memory string for the protobuf

~/anaconda3/envs/amp/lib/python3.7/site-packages/torch/onnx/utils.py in _model_to_graph(model, args, verbose, training, input_names, output_names, operator_export_type, example_outputs, propagate, _retain_param_name, do_constant_folding, _disable_torch_constant_prop, fixed_batch_size)
    277             model.graph, tuple(in_vars), False, propagate)
    278     else:
--> 279         graph, torch_out = _trace_and_get_graph_from_model(model, args, training)
    280         state_dict = _unique_state_dict(model)
    281         params = list(state_dict.values())

~/anaconda3/envs/amp/lib/python3.7/site-packages/torch/onnx/utils.py in _trace_and_get_graph_from_model(model, args, training)
    234     # training mode was.)
    235     with set_training(model, training):
--> 236         trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph(model, args, _force_outplace=True, _return_inputs_states=True)
    237         warn_on_static_input_change(inputs_states)
    238 

~/anaconda3/envs/amp/lib/python3.7/site-packages/torch/jit/__init__.py in _get_trace_graph(f, args, kwargs, _force_outplace, return_inputs, _return_inputs_states)
    275     if not isinstance(args, tuple):
    276         args = (args,)
--> 277     outs = ONNXTracedModule(f, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
    278     return outs
    279 

~/anaconda3/envs/amp/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    530             result = self._slow_forward(*input, **kwargs)
    531         else:
--> 532             result = self.forward(*input, **kwargs)
    533         for hook in self._forward_hooks.values():
    534             hook_result = hook(self, input, result)

~/anaconda3/envs/amp/lib/python3.7/site-packages/torch/jit/__init__.py in forward(self, *args)
    358             in_vars + module_state,
    359             _create_interpreter_name_lookup_fn(),
--> 360             self._force_outplace,
    361         )
    362 

~/anaconda3/envs/amp/lib/python3.7/site-packages/torch/jit/__init__.py in wrapper(*args)
    345             if self._return_inputs_states:
    346                 inputs_states.append(_unflatten(args[:len(in_vars)], in_desc))
--> 347             outs.append(self.inner(*trace_inputs))
    348             if self._return_inputs_states:
    349                 inputs_states[0] = (inputs_states[0], trace_inputs)

~/anaconda3/envs/amp/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    528                 input = result
    529         if torch._C._get_tracing_state():
--> 530             result = self._slow_forward(*input, **kwargs)
    531         else:
    532             result = self.forward(*input, **kwargs)

~/anaconda3/envs/amp/lib/python3.7/site-packages/torch/nn/modules/module.py in _slow_forward(self, *input, **kwargs)
    514                 recording_scopes = False
    515         try:
--> 516             result = self.forward(*input, **kwargs)
    517         finally:
    518             if recording_scopes:

<ipython-input-2-a07de6cca1a4> in forward(self, x)
      5 
      6     def forward(self, x):
----> 7         return self.layer(x)

~/anaconda3/envs/amp/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    528                 input = result
    529         if torch._C._get_tracing_state():
--> 530             result = self._slow_forward(*input, **kwargs)
    531         else:
    532             result = self.forward(*input, **kwargs)

~/anaconda3/envs/amp/lib/python3.7/site-packages/torch/nn/modules/module.py in _slow_forward(self, *input, **kwargs)
    514                 recording_scopes = False
    515         try:
--> 516             result = self.forward(*input, **kwargs)
    517         finally:
    518             if recording_scopes:

~/anaconda3/envs/amp/lib/python3.7/site-packages/torch/nn/modules/conv.py in forward(self, input)
    343 
    344     def forward(self, input):
--> 345         return self.conv2d_forward(input, self.weight)
    346 
    347 class Conv3d(_ConvNd):

~/anaconda3/envs/amp/lib/python3.7/site-packages/torch/nn/modules/conv.py in conv2d_forward(self, input, weight)
    340                             _pair(0), self.dilation, self.groups)
    341         return F.conv2d(input, weight, self.bias, self.stride,
--> 342                         self.padding, self.dilation, self.groups)
    343 
    344     def forward(self, input):

~/anaconda3/envs/amp/lib/python3.7/site-packages/apex/amp/wrap.py in wrapper(*args, **kwargs)
     26                                      args,
     27                                      kwargs)
---> 28         return orig_fn(*new_args, **kwargs)
     29     return wrapper
     30 

RuntimeError: Cannot insert a Tensor that requires grad as a constant. Consider making it a parameter or input, or detaching the gradient
Tensor:
(1,1,.,.) = 
  0.0913 -0.0682 -0.0854 -0.1993 -0.1993
 -0.0832  0.0637  0.1675  0.1593 -0.0156
 -0.1478  0.1442  0.0658  0.1543  0.0630
  0.0398  0.1490 -0.1765  0.0949  0.1530
 -0.1447  0.0278 -0.1873  0.0504 -0.0494
[ torch.cuda.HalfTensor{1,1,5,5} ]```

Reproduce the failure using the Foo module, the apex amp.initialize call, the preceding forward pass, and torch.onnx.export shown in the report. Start at the torch.onnx.export tracing path and the apex-wrapped Conv2d call; done means the AMP-initialized model exports successfully after a forward pass without the reported constant-gradient error.