diff --git a/py/torch_tensorrt/csrc/torch_tensorrt_py.cpp b/py/torch_tensorrt/csrc/torch_tensorrt_py.cpp index 6b1ffd4ccf..4b87dea604 100644 --- a/py/torch_tensorrt/csrc/torch_tensorrt_py.cpp +++ b/py/torch_tensorrt/csrc/torch_tensorrt_py.cpp @@ -23,7 +23,14 @@ class pyCalibratorTrampoline : public Derived { using Derived::Derived; // Inherit constructors int getBatchSize() const noexcept override { - PYBIND11_OVERLOAD_PURE_NAME(int, Derived, "get_batch_size", getBatchSize); + try { + PYBIND11_OVERLOAD_PURE_NAME(int, Derived, "get_batch_size", getBatchSize); + } catch (std::exception const& e) { + LOG_ERROR("Exception caught in get_batch_size" + std::string(e.what())); + } catch (...) { + LOG_ERROR("Exception caught in get_batch_size"); + } + return -1; } bool getBatch(void* bindings[], const char* names[], int nbBindings) noexcept override { @@ -71,8 +78,15 @@ class pyIInt8Calibrator : public pyCalibratorTrampoline