Repository: pytorch
Issue: Transform caffe2 to trt failed
Labels: caffe2,caffe2-op,triaged

## 🐛 Bug
I called transform_caffe2_net(pred_net,
                                            {input_name: input_blob_dims},
                                            max_batch_size=1, debug_builder=False,
                                            build_serializable_op=False)

to transform caffe2 model to trt, got the following exception, input_blob_dims = (1, 3, 800, 832), input_name = pred_net.external_input[0]:

Traceback (most recent call last):
  File "test.py", line 12, in <module>
    ir = image_recognition_Casrcnn.ImgRecog_detectron(0)
  File "/work/test_code/image_recognition_Casrcnn.py", line 141, in __init__
    build_serializable_op=False)
  File "/workspace/caffe2-pytorch/build/caffe2/python/trt/transform.py", line 107, in transform_caffe2_net
    build_serializable_op)
RuntimeError: Expected axis_index < ndims to be true, but got false.  (Could this error message be improved?  If so, please report an enhancement request to PyTorch.) (canonical_axis_index_ at /workspace/caffe2-pytorch/c10/core/TensorImpl.h:92)
frame #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&) + 0x6c (0x7fb51c7fbf4c in /workspace/caffe2-pytorch/build/lib/libc10.so)
frame #1: <unknown function> + 0x29caf44 (0x7fb4de161f44 in /workspace/caffe2-pytorch/build/lib/libtorch.so)
frame #2: caffe2::onnx::OnnxExporter::CreateConcatNodes(caffe2::OperatorDef const&, std::unordered_map<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, caffe2::TensorShape, std::hash<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > >, std::equal_to<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > >, std::allocator<std::pair<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const, caffe2::TensorShape> > > const&) + 0xb1f (0x7fb4de2da9ef in /workspace/caffe2-pytorch/build/lib/libtorch.so)
frame #3: caffe2::onnx::OnnxExporter::Caffe2OpToOnnxNodes(caffe2::OperatorDef const&, std::unordered_map<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, caffe2::TensorShape, std::hash<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > >, std::equal_to<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > >, std::allocator<std::pair<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const, caffe2::TensorShape> > > const&) + 0xd6 (0x7fb4de2cd9d6 in /workspace/caffe2-pytorch/build/lib/libtorch.so)
frame #4: <unknown function> + 0x4288ec2 (0x7fb4dfa1fec2 in /workspace/caffe2-pytorch/build/lib/libtorch.so)
frame #5: caffe2::opt::OptimizeForBackend(caffe2::NetDef&, std::function<bool (caffe2::OperatorDef const&)>, std::function<caffe2::NetDef (caffe2::NetDef const&)>, bool) + 0x5ee (0x7fb4de887ace in /workspace/caffe2-pytorch/build/lib/libtorch.so)
frame #6: caffe2::TensorRTTransformer::Transform(caffe2::Workspace*, caffe2::NetDef*, std::unordered_map<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, caffe2::TensorShape, std::hash<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > >, std::equal_to<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > >, std::allocator<std::pair<std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const, caffe2::TensorShape> > > const&) + 0x666 (0x7fb4dfa26da6 in /workspace/caffe2-pytorch/build/lib/libtorch.so)
frame #7: <unknown function> + 0x105e50 (0x7fb51cd4de50 in /workspace/caffe2-pytorch/build/caffe2/python/caffe2_pybind11_state_gpu.so)
frame #8: <unknown function> + 0x9af8d (0x7fb51cce2f8d in /workspace/caffe2-pytorch/build/caffe2/python/caffe2_pybind11_state_gpu.so)
<omitting python frames>
frame #27: __libc_start_main + 0xe7 (0x7fb5283c3b97 in /lib/x86_64-linux-gnu/libc.so.6)