MaskRCNN from ScriptModule to ONNX - Unknown Type BoxCoder

🐛 Describe the bug

While attempting to create an ONNX version of Maskrcnn, starting from a ScriptModule, an error occurs, indicating that __torch__.torchvision.models.detection._utils.BoxCoder is an unknown type.

MWE:

import torch
from torchvision.models.detection.mask_rcnn import maskrcnn_resnet50_fpn

model = maskrcnn_resnet50_fpn()
model.eval()
script_model = torch.jit.script(model)
example_image = torch.rand((3, 800, 1000))
torch.onnx.export(
	script_model,
	[example_image],
	"test.onnx",
	example_outputs=script_model([example_image])[1], # index 0 is losses
	opset_version = 11
)

Error traceback:

Traceback (most recent call last):
  File "/home/nmota/test_onnx.py", line 8, in <module>
    torch.onnx.export(
  File "/usr/lib/python3.9/site-packages/torch/onnx/__init__.py", line 275, in export
    return utils.export(model, args, f, export_params, verbose, training,
  File "/usr/lib/python3.9/site-packages/torch/onnx/utils.py", line 88, in export
    _export(model, args, f, export_params, verbose, training, input_names, output_names,
  File "/usr/lib/python3.9/site-packages/torch/onnx/utils.py", line 689, in _export
    _model_to_graph(model, args, verbose, input_names,
  File "/usr/lib/python3.9/site-packages/torch/onnx/utils.py", line 458, in _model_to_graph
    graph, params, torch_out, module = _create_jit_graph(model, args,
  File "/usr/lib/python3.9/site-packages/torch/onnx/utils.py", line 402, in _create_jit_graph
    module, params = torch._C._jit_onnx_list_model_parameters(freezed_m)
RuntimeError: 
Unknown type __torch__.torchvision.models.detection._utils.BoxCoder (of Python compilation unit at: 0x55bbdb787f00) encountered in handling model params. This class type does not extend __getstate__ method.:

✗ - status code 1

Unfortunately, I cannot test with a more recent version. Is this something that has been fixed recently?

Versions

PyTorch version: 1.9.0
Is debug build: False
CUDA used to build PyTorch: 11.4
ROCM used to build PyTorch: N/A

OS: Arch Linux (x86_64)
GCC version: (GCC) 11.1.0
Clang version: 12.0.1
CMake version: version 3.21.2
Libc version: glibc-2.33

Python version: 3.9.6 (default, Jun 30 2021, 10:22:16)  [GCC 11.1.0] (64-bit runtime)
Python platform: Linux-5.13.13-arch1-1-x86_64-with-glibc2.33
Is CUDA available: True
CUDA runtime version: 11.4.100
GPU models and configuration: 
GPU 0: NVIDIA TITAN X (Pascal)

Nvidia driver version: 470.63.01
cuDNN version: Probably one of the following:
/usr/lib/libcudnn.so.8.2.2
/usr/lib/libcudnn_adv_infer.so.8.2.2
/usr/lib/libcudnn_adv_train.so.8.2.2
/usr/lib/libcudnn_cnn_infer.so.8.2.2
/usr/lib/libcudnn_cnn_train.so.8.2.2
/usr/lib/libcudnn_ops_infer.so.8.2.2
/usr/lib/libcudnn_ops_train.so.8.2.2
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.20.3
[pip3] torch==1.9.0
[pip3] torchvision==0.10.0a0
[conda] Could not collect

cc @neginraoof

Start with the MaskRCNN MWE using maskrcnn_resnet50_fpn, torch.jit.script, and torch.onnx.export, then inspect how the detection model's BoxCoder is handled during ONNX parameter processing. Reproduce with the reported PyTorch 1.9.0 and torchvision 0.10.0a0 versions if available, and check whether export completes without the unknown-type error.