[ONNX] Update embedding_bag export (#44693)

Summary:
Export of embedding bag with dynamic list of offsets.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/44693

Reviewed By: malfet

Differential Revision: D23831980

Pulled By: bzinodev

fbshipit-source-id: 3eaff1a0f20d1bcfb8039e518d78c491be381e1a
This commit is contained in:
Negin Raoof
2020-09-30 13:34:26 -07:00
committed by Facebook GitHub Bot
parent ac9a708ed0
commit 6b42ca2d69
10 changed files with 218 additions and 111 deletions

View File

@ -452,6 +452,7 @@ void initJITBindings(PyObject* module) {
})
.def("_jit_pass_onnx_block", BlockToONNX)
.def("_jit_pass_fixup_onnx_controlflow_node", FixupONNXControlflowNode)
.def("_jit_pass_fixup_onnx_loop_node_inputs", FixupONNXLoopNodeInputs)
.def("_jit_pass_canonicalize_graph_fuser_ops", CanonicalizeOps)
.def("_jit_pass_decompose_ops", DecomposeOps)
.def("_jit_pass_specialize_autogradzero", specializeAutogradZero)