[caffe2] RuntimeError: block->inputs().size() >= num_initializers ASSERT FAILED

I am using the caffe2 to convert the pytorch model that trained on pytorch 0.4.1. After installed, my pytorch version is

torch.__version__              = 1.0.0.dev20181125
torch.version.cuda             = 9.2.148
torch.backends.cudnn.version() = 7401

My model is trained on pytorch 0.4.1 using conv3d,conv3d_transpode, max_pooling 3d, Batchnorm3d and Sequential as Densenet example. I used the code to convert to onnx model

model.load_state_dict('checkpoint.pth')
input = torch.randn(1, 1, 32, 32, 32, requires_grad=True)
input = input.to(device)
# Export the model
torch_out = torch.onnx._export(model,
					   input,                      
					   "model_caff2.onnx", 
					   export_params=True)     

However, I got the error as bellow, how should I fix it? This is the full log

Traceback (most recent call last):
  File "convert_pytorch2caffe2.py", line 60, in <module>
    export_params=True)      # store the trained parameter weights inside the model file
  File "/home/john/anaconda3/lib/python3.6/site-packages/torch/onnx/__init__.py", line 22, in _export
    return utils._export(*args, **kwargs)
  File "/home/john/anaconda3/lib/python3.6/site-packages/torch/onnx/utils.py", line 287, in _export
    proto, export_map = graph.export(params, _onnx_opset_version, defer_weight_export, operator_export_type)
RuntimeError: block->inputs().size() >= num_initializers ASSERT FAILED at /opt/conda/conda-bld/pytorch-nightly_1543136687659/work/torch/csrc/jit/export.cpp:299, please report a bug to PyTorch. (EncodeBlock at /opt/conda/conda-bld/pytorch-nightly_1543136687659/work/torch/csrc/jit/export.cpp:299)
frame #0: c10::Error::Error(c10::SourceLocation, std::string const&) + 0x45 (0x7f36392ab5e5 in /home/john/anaconda3/lib/python3.6/site-packages/torch/lib/libc10.so)
frame #1: <unknown function> + 0x5380aa (0x7f36369270aa in /home/john/anaconda3/lib/python3.6/site-packages/torch/lib/libtorch.so.1)
frame #2: <unknown function> + 0x53b95e (0x7f363692a95e in /home/john/anaconda3/lib/python3.6/site-packages/torch/lib/libtorch.so.1)
frame #3: torch::jit::ExportGraph(std::shared_ptr<torch::jit::Graph> const&, std::vector<at::Tensor, std::allocator<at::Tensor> > const&, long, bool, torch::onnx::OperatorExportTypes) + 0x37 (0x7f363692aab7 in /home/john/anaconda3/lib/python3.6/site-packages/torch/lib/libtorch.so.1)
frame #4: <unknown function> + 0x402988 (0x7f36699c8988 in /home/john/anaconda3/lib/python3.6/site-packages/torch/lib/libtorch_python.so)
frame #5: <unknown function> + 0x191206 (0x7f3669757206 in /home/john/anaconda3/lib/python3.6/site-packages/torch/lib/libtorch_python.so)
<omitting python frames>
frame #24: __libc_start_main + 0xf0 (0x7f367cfdd830 in /lib/x86_64-linux-gnu/libc.so.6)