feat(//core/lowering): additional logging in module fallback

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

core/lowering/BUILD

@@ -13,6 +13,7 @@ cc_library(
         "drop_unused_nodes.cpp",
         "lowering.cpp",
         "register_trt_placeholder_ops.cpp",
+        "LowerInfo.cpp"
     ],
     hdrs = [
         "lowering.h",

core/lowering/LowerInfo.cpp

@@ -0,0 +1,23 @@
+#include <iostream>
+#include <sstream>
+#include <utility>
+
+#include "core/lowering/lowering.h"
+
+namespace trtorch {
+namespace core {
+namespace lowering {
+
+std::ostream& operator<<(std::ostream& os, const LowerInfo& l) {
+    os << "Settings requested for Lowering:" << std::endl;
+    os << "    Forced Fallback Modules: [" << std::endl;
+    for (auto i : l.forced_fallback_modules) {
+        os << "      " << i << std::endl;
+    }
+    os << "    ]";
+    return os;
+}
+
+} // namespace lowering
+} // namespace core
+} // namespace trtorch

py/trtorch/_compiler.py

@@ -59,6 +59,9 @@ def compile(module: torch.jit.ScriptModule, compile_spec: Any) -> torch.jit.Scri
                         "force_fallback_ops": [
                             "aten::max_pool2d" # List of specific ops to require running in PyTorch
                         ],
+                        "force_fallback_modules": [
+                            "mypymod.mytorchmod" # List of specific torch modules to require running in PyTorch
+                        ],
                         "min_block_size": 3 # Minimum number of ops an engine must incapsulate to be run in TensorRT
                     }
                 }