fix: register the torch_fallback attribute in Python API

Signed-off-by: Bo Wang <wangbo1995ee@163.com>
pytorch · Mar 11, 2021 · 8b7919f · 8b7919f
1 parent 100b090
commit 8b7919f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/py/trtorch/csrc/trtorch_py.cpp b/py/trtorch/csrc/trtorch_py.cpp
@@ -115,7 +115,8 @@ PYBIND11_MODULE(_C, m) {
       .def_readwrite("num_min_timing_iters", &CompileSpec::num_min_timing_iters)
       .def_readwrite("num_avg_timing_iters", &CompileSpec::num_avg_timing_iters)
       .def_readwrite("workspace_size", &CompileSpec::workspace_size)
-      .def_readwrite("max_batch_size", &CompileSpec::max_batch_size);
+      .def_readwrite("max_batch_size", &CompileSpec::max_batch_size)
+      .def_readwrite("torch_fallback", &CompileSpec::torch_fallback);
 
   py::class_<Device>(m, "Device")
       .def(py::init<>())