diff --git a/core/compiler.cpp b/core/compiler.cpp
index cf4d96668c..29a5e0c510 100644
--- a/core/compiler.cpp
+++ b/core/compiler.cpp
@@ -214,7 +214,7 @@ torch::jit::script::Module CompileGraphWithFallback(const torch::jit::script::Mo
 
       // segment the graph and convert segmented TensorRT block
       auto segmented_blocks =
-          partitioning::Partition(g, convert_cfg.input_ranges, convert_cfg.engine_settings.torch_fallback);
+          partitioning::Partition(g, convert_cfg.input_ranges, cfg.partition_info);
       if (segmented_blocks.size() == 1 && segmented_blocks[0].target() == partitioning::SegmentedBlock::kTorch) {
         return mod;
       }
@@ -223,9 +223,9 @@ torch::jit::script::Module CompileGraphWithFallback(const torch::jit::script::Mo
       std::unordered_map<torch::jit::Value*, torch::jit::Value*> old_to_new_g;
       for (auto& seg_block : segmented_blocks) {
         if (seg_block.target() == partitioning::SegmentedBlock::kTensorRT) {
-          std::vector<conversion::InputRange> input_ranges;
+          std::vector<ir::InputRange> input_ranges;
           for (auto& shape : seg_block.in_shape()) {
-            input_ranges.push_back(conversion::InputRange(util::toVec(shape)));
+            input_ranges.push_back(ir::InputRange(util::toVec(shape)));
           }
           // update the input ranges for each segments
           convert_cfg.input_ranges = input_ranges;
@@ -258,7 +258,7 @@ torch::jit::script::Module CompileGraphWithFallback(const torch::jit::script::Mo
 
 torch::jit::script::Module CompileGraph(const torch::jit::script::Module& mod, CompileSpec cfg) {
   // TODO: not sure how to deal with duplicated code here, so just cut out a branch temporally
-  if (cfg.convert_info.engine_settings.torch_fallback.enabled) {
+  if (cfg.partition_info.enabled) {
     return CompileGraphWithFallback(mod, cfg);
   }
   // TODO: Should be doing a functional transform but need PR #31978
diff --git a/core/compiler.h b/core/compiler.h
index 512b30123f..403d5e4e5b 100644
--- a/core/compiler.h
+++ b/core/compiler.h
@@ -3,14 +3,17 @@
 #include <cuda_runtime.h>
 #include <vector>
 #include "core/conversion/conversion.h"
+#include "core/ir/ir.h"
+#include "core/partitioning/partitioning.h"
 #include "torch/csrc/jit/api/module.h"
 
 namespace trtorch {
 namespace core {
 
 struct CompileSpec {
-  CompileSpec(std::vector<conversion::InputRange> input_ranges) : convert_info(std::move(input_ranges)) {}
+  CompileSpec(std::vector<ir::InputRange> input_ranges) : convert_info(std::move(input_ranges)) {}
   conversion::ConversionInfo convert_info;
+  partitioning::PartitionInfo partition_info;
 };
 
 bool CheckMethodOperatorSupport(const torch::jit::script::Module& mod, std::string method_name);
diff --git a/core/conversion/BUILD b/core/conversion/BUILD
index 22fc9bd189..dcce0758eb 100644
--- a/core/conversion/BUILD
+++ b/core/conversion/BUILD
@@ -23,7 +23,8 @@ cc_library(
         "//core/conversion/conversionctx",
         "//core/conversion/converters",
         "//core/conversion/evaluators",
-        "//core/util:prelude"
+        "//core/util:prelude",
+        "//core/ir",
     ] + select({
         ":use_pre_cxx11_abi":  ["@libtorch_pre_cxx11_abi//:libtorch"],
         "//conditions:default":  ["@libtorch//:libtorch"],
diff --git a/core/conversion/InterfaceTypes.cpp b/core/conversion/InterfaceTypes.cpp
index 14853c92e3..f7d6565f26 100644
--- a/core/conversion/InterfaceTypes.cpp
+++ b/core/conversion/InterfaceTypes.cpp
@@ -23,55 +23,6 @@ GraphParams get_named_params(c10::ArrayRef<torch::jit::Value*> inputs, std::vect
   return std::move(named_params);
 }
 
-InputRange::InputRange(std::vector<int64_t> d) {
-  if (d.size() > 5) {
-    LOG_WARNING("Verify that this dim size is accepted");
-  }
-
-  opt = util::toDims(d);
-  min = util::toDims(d);
-  max = util::toDims(d);
-  input_shape = util::toDims(d);
-  input_is_dynamic = false;
-}
-
-InputRange::InputRange(std::vector<int64_t> min_shape, std::vector<int64_t> opt_shape, std::vector<int64_t> max_shape) {
-  if (min_shape.size() > 5 || opt_shape.size() > 5 || max_shape.size() > 5) {
-    LOG_WARNING("Verify that this dim size is accepted");
-  }
-
-  std::set<size_t> sizes;
-  sizes.insert(min_shape.size());
-  sizes.insert(opt_shape.size());
-  sizes.insert(max_shape.size());
-
-  if (sizes.size() != 1) {
-    LOG_ERROR(
-        "Expected all input sizes have the same dimensions, but found dimensions: min("
-        << min_shape.size() << "), opt(" << opt_shape.size() << "), max(" << max_shape.size() << ")");
-  }
-
-  min = util::toDims(min_shape);
-  opt = util::toDims(opt_shape);
-  max = util::toDims(max_shape);
-
-  std::vector<int64_t> dyn_shape;
-  for (size_t i = 0; i < opt_shape.size(); i++) {
-    std::set<uint64_t> dim;
-    dim.insert(min_shape[i]);
-    dim.insert(opt_shape[i]);
-    dim.insert(max_shape[i]);
-    if (dim.size() != 1) {
-      dyn_shape.push_back(-1);
-      input_is_dynamic = true;
-    } else {
-      dyn_shape.push_back(opt_shape[i]);
-    }
-  }
-
-  input_shape = util::toDims(dyn_shape);
-}
-
 } // namespace conversion
 } // namespace core
 } // namespace trtorch
diff --git a/core/conversion/conversion.cpp b/core/conversion/conversion.cpp
index fb42a41dbf..6d5393b91d 100644
--- a/core/conversion/conversion.cpp
+++ b/core/conversion/conversion.cpp
@@ -118,7 +118,7 @@ void AddLayer(ConversionCtx* ctx, const torch::jit::Node* n) {
                        << "please report this error to https://www.github.com/NVIDIA/TRTorch/issues");
 }
 
-void AddInputs(ConversionCtx* ctx, at::ArrayRef<const torch::jit::Value*> inputs, std::vector<InputRange>& input_dims) {
+void AddInputs(ConversionCtx* ctx, at::ArrayRef<const torch::jit::Value*> inputs, std::vector<ir::InputRange>& input_dims) {
   std::vector<const torch::jit::Value*> input_tensors;
   for (auto in : inputs) {
     // Disregarding inputs that are not tensors
diff --git a/core/conversion/conversion.h b/core/conversion/conversion.h
index b91e769932..f4380c72ec 100644
--- a/core/conversion/conversion.h
+++ b/core/conversion/conversion.h
@@ -4,27 +4,17 @@
 
 #include "NvInfer.h"
 #include "core/conversion/conversionctx/ConversionCtx.h"
+#include "core/ir/ir.h"
 #include "torch/csrc/jit/ir/ir.h"
 
 namespace trtorch {
 namespace core {
 namespace conversion {
 
-struct InputRange {
-  nvinfer1::Dims min;
-  nvinfer1::Dims max;
-  nvinfer1::Dims opt;
-  nvinfer1::Dims input_shape;
-  bool input_is_dynamic = false;
-  // Should we restrict to unsigned?
-  InputRange(std::vector<int64_t> d);
-  InputRange(std::vector<int64_t> min_shape, std::vector<int64_t> opt_shape, std::vector<int64_t> max_shape);
-};
-
 struct ConversionInfo {
-  std::vector<InputRange> input_ranges;
+  std::vector<ir::InputRange> input_ranges;
   BuilderSettings engine_settings;
-  ConversionInfo(std::vector<InputRange> input_ranges)
+  ConversionInfo(std::vector<ir::InputRange> input_ranges)
       : input_ranges(std::move(input_ranges)), engine_settings(BuilderSettings()) {}
 };
 
diff --git a/core/conversion/conversionctx/ConversionCtx.cpp b/core/conversion/conversionctx/ConversionCtx.cpp
index 97c7606bdb..04f6aafe5c 100644
--- a/core/conversion/conversionctx/ConversionCtx.cpp
+++ b/core/conversion/conversionctx/ConversionCtx.cpp
@@ -36,17 +36,6 @@ std::ostream& operator<<(std::ostream& os, const BuilderSettings& s) {
     }
     os << "\n    Engine Capability: " << s.capability                                      \
        << "\n    Calibrator Created: " << (s.calibrator != nullptr);
-
-    os << "\n    Torch Fallback: " << s.torch_fallback.enabled;
-    if (s.torch_fallback.enabled) {
-      os << "\n    Fallback Min Block Size: " << s.torch_fallback.min_block_size;
-      if (!s.torch_fallback.forced_fallback_operators.empty()) {
-        os << "\n    Forced Fallback Operators:";
-        for (auto it = s.torch_fallback.forced_fallback_operators.begin(); it != s.torch_fallback.forced_fallback_operators.end(); ++it) {
-          os << " " << *it;
-        }
-      }
-    }
     return os;
 }
 // clang-format on
diff --git a/core/conversion/conversionctx/ConversionCtx.h b/core/conversion/conversionctx/ConversionCtx.h
index 95cfebd874..e191d84512 100644
--- a/core/conversion/conversionctx/ConversionCtx.h
+++ b/core/conversion/conversionctx/ConversionCtx.h
@@ -22,12 +22,6 @@ struct Device {
   Device() : device_type(nvinfer1::DeviceType::kGPU), gpu_id(0), dla_core(0), allow_gpu_fallback(false) {}
 };
 
-struct TorchFallback {
-  bool enabled = false;
-  uint64_t min_block_size = 1;
-  std::vector<std::string> forced_fallback_operators;
-};
-
 struct BuilderSettings {
   nvinfer1::DataType op_precision = nvinfer1::DataType::kFLOAT;
   bool disable_tf32 = false;
@@ -36,7 +30,6 @@ struct BuilderSettings {
   bool strict_types = false;
   bool truncate_long_and_double = false;
   Device device;
-  TorchFallback torch_fallback;
   nvinfer1::EngineCapability capability = nvinfer1::EngineCapability::kDEFAULT;
   nvinfer1::IInt8Calibrator* calibrator = nullptr;
   uint64_t num_min_timing_iters = 2;
diff --git a/core/ir/BUILD b/core/ir/BUILD
new file mode 100644
index 0000000000..e889337608
--- /dev/null
+++ b/core/ir/BUILD
@@ -0,0 +1,35 @@
+package(default_visibility = ["//visibility:public"])
+
+config_setting(
+    name = "use_pre_cxx11_abi",
+    values = {
+        "define": "abi=pre_cxx11_abi",
+    }
+)
+
+cc_library(
+    name = "ir",
+    hdrs = [
+        "ir.h"
+    ],
+    srcs = [
+        "InputRange.cpp",
+    ],
+    deps = [
+        "@tensorrt//:nvinfer",
+        "//core/util:prelude",
+    ] + select({
+        ":use_pre_cxx11_abi":  ["@libtorch_pre_cxx11_abi//:libtorch"],
+        "//conditions:default":  ["@libtorch//:libtorch"],
+    }),
+)
+
+load("@rules_pkg//:pkg.bzl", "pkg_tar")
+
+pkg_tar(
+    name = "include",
+    package_dir = "core/ir/",
+    srcs = [
+        "ir.h",
+    ],
+)
diff --git a/core/ir/InputRange.cpp b/core/ir/InputRange.cpp
new file mode 100644
index 0000000000..8c0adda8dc
--- /dev/null
+++ b/core/ir/InputRange.cpp
@@ -0,0 +1,59 @@
+#include "core/ir/ir.h"
+#include "core/util/prelude.h"
+
+namespace trtorch {
+namespace core {
+namespace ir {
+
+InputRange::InputRange(std::vector<int64_t> d) {
+  if (d.size() > 5) {
+    LOG_WARNING("Verify that this dim size is accepted");
+  }
+
+  opt = util::toDims(d);
+  min = util::toDims(d);
+  max = util::toDims(d);
+  input_shape = util::toDims(d);
+  input_is_dynamic = false;
+}
+
+InputRange::InputRange(std::vector<int64_t> min_shape, std::vector<int64_t> opt_shape, std::vector<int64_t> max_shape) {
+  if (min_shape.size() > 5 || opt_shape.size() > 5 || max_shape.size() > 5) {
+    LOG_WARNING("Verify that this dim size is accepted");
+  }
+
+  std::set<size_t> sizes;
+  sizes.insert(min_shape.size());
+  sizes.insert(opt_shape.size());
+  sizes.insert(max_shape.size());
+
+  if (sizes.size() != 1) {
+    LOG_ERROR(
+        "Expected all input sizes have the same dimensions, but found dimensions: min("
+        << min_shape.size() << "), opt(" << opt_shape.size() << "), max(" << max_shape.size() << ")");
+  }
+
+  min = util::toDims(min_shape);
+  opt = util::toDims(opt_shape);
+  max = util::toDims(max_shape);
+
+  std::vector<int64_t> dyn_shape;
+  for (size_t i = 0; i < opt_shape.size(); i++) {
+    std::set<uint64_t> dim;
+    dim.insert(min_shape[i]);
+    dim.insert(opt_shape[i]);
+    dim.insert(max_shape[i]);
+    if (dim.size() != 1) {
+      dyn_shape.push_back(-1);
+      input_is_dynamic = true;
+    } else {
+      dyn_shape.push_back(opt_shape[i]);
+    }
+  }
+
+  input_shape = util::toDims(dyn_shape);
+}
+
+} // namespace ir
+} // namespace core
+} // namespace trtorch
\ No newline at end of file
diff --git a/core/ir/ir.h b/core/ir/ir.h
new file mode 100644
index 0000000000..87bae52645
--- /dev/null
+++ b/core/ir/ir.h
@@ -0,0 +1,23 @@
+#pragma once
+
+#include <vector>
+#include "NvInfer.h"
+
+namespace trtorch {
+namespace core {
+namespace ir {
+
+struct InputRange {
+  nvinfer1::Dims min;
+  nvinfer1::Dims max;
+  nvinfer1::Dims opt;
+  nvinfer1::Dims input_shape;
+  bool input_is_dynamic = false;
+  // Should we restrict to unsigned?
+  InputRange(std::vector<int64_t> d);
+  InputRange(std::vector<int64_t> min_shape, std::vector<int64_t> opt_shape, std::vector<int64_t> max_shape);
+};
+
+} // namespace ir
+} // namespace core
+} // namespace trtorch
\ No newline at end of file
diff --git a/core/partitioning/BUILD b/core/partitioning/BUILD
index 0f21667718..85e7dc3d8c 100644
--- a/core/partitioning/BUILD
+++ b/core/partitioning/BUILD
@@ -12,6 +12,7 @@ cc_library(
     hdrs = [
         "SegmentedBlock.h",
         "shape_analysis.h",
+        "PartitionInfo.h",
         "partitioning.h",
     ],
     srcs = [
@@ -20,8 +21,9 @@ cc_library(
         "partitioning.cpp",
     ],
     deps = [
-        "//core/conversion",
         "//core/util:prelude",
+        "//core/ir",
+        "//core/conversion",
         "//core/lowering"
     ] + select({
         ":use_pre_cxx11_abi":  ["@libtorch_pre_cxx11_abi//:libtorch"],
@@ -35,6 +37,11 @@ load("@rules_pkg//:pkg.bzl", "pkg_tar")
 pkg_tar(
     name = "include",
     package_dir = "core/partitioning/",
-    srcs = ["partitioning.h"],
+    srcs = [
+        "SegmentedBlock.h",
+        "shape_analysis.h",
+        "PartitionInfo.h",
+        "partitioning.h",
+    ],
 )
 
diff --git a/core/partitioning/PartitionInfo.h b/core/partitioning/PartitionInfo.h
new file mode 100644
index 0000000000..d1a8aca321
--- /dev/null
+++ b/core/partitioning/PartitionInfo.h
@@ -0,0 +1,19 @@
+#pragma once
+
+#include <cstdint>
+#include <vector>
+#include <string>
+
+namespace trtorch {
+namespace core {
+namespace partitioning {
+
+struct PartitionInfo {
+    bool enabled = false;
+    uint64_t min_block_size = 1;
+    std::vector<std::string> forced_fallback_operators;
+};
+
+} // namespace partitioning
+} // namespace core
+} // namespace trtorch
\ No newline at end of file
diff --git a/core/partitioning/SegmentedBlock.cpp b/core/partitioning/SegmentedBlock.cpp
index 96a765a25f..1bbf25197e 100644
--- a/core/partitioning/SegmentedBlock.cpp
+++ b/core/partitioning/SegmentedBlock.cpp
@@ -38,51 +38,7 @@ torch::jit::Node* SegmentedBlock::cloneNode(torch::jit::Node* node) {
   return new_node;
 }
 
-std::vector<SegmentedBlock> segment_graph(
-    std::shared_ptr<torch::jit::Graph> g,
-    const conversion::TorchFallback& fallback_info) {
-  auto min_block_size = fallback_info.min_block_size;
-  std::unordered_set<std::string> forced_fallback_operators(
-      fallback_info.forced_fallback_operators.begin(), fallback_info.forced_fallback_operators.end());
 
-  auto nodes = g->block()->nodes();
-  std::vector<SegmentedBlock> segmented_blocks;
-
-  // segment the nodes
-  std::vector<torch::jit::Node*> tensorrt_nodes, pytorch_nodes;
-  for (const auto n : nodes) {
-    if (n->kind() == torch::jit::prim::Constant)
-      continue;
-
-    std::string node_string(n->kind().toQualString());
-    if (conversion::OpSupported(n) && !forced_fallback_operators.count(node_string)) {
-      tensorrt_nodes.push_back(n);
-      if (tensorrt_nodes.size() >= min_block_size && !pytorch_nodes.empty()) {
-        segmented_blocks.emplace_back(SegmentedBlock::kTorch, pytorch_nodes);
-        pytorch_nodes.clear();
-      }
-    } else {
-      if (tensorrt_nodes.size() >= min_block_size) {
-        segmented_blocks.emplace_back(SegmentedBlock::kTensorRT, tensorrt_nodes);
-      } else {
-        pytorch_nodes.insert(pytorch_nodes.end(), tensorrt_nodes.begin(), tensorrt_nodes.end());
-      }
-      tensorrt_nodes.clear();
-      pytorch_nodes.push_back(n);
-    }
-  }
-
-  // if there is any kTorch nodes left, then either the last nodes are kTorch or last nodes are kTensorRT but num <
-  // min_block_size
-  if (!pytorch_nodes.empty()) {
-    pytorch_nodes.insert(pytorch_nodes.end(), tensorrt_nodes.begin(), tensorrt_nodes.end());
-    segmented_blocks.emplace_back(SegmentedBlock::kTorch, pytorch_nodes);
-  } else {
-    segmented_blocks.emplace_back(SegmentedBlock::kTensorRT, tensorrt_nodes);
-  }
-
-  return std::move(segmented_blocks);
-}
 
 } // namespace partitioning
 } // namespace core
diff --git a/core/partitioning/SegmentedBlock.h b/core/partitioning/SegmentedBlock.h
index b9a9c74ea9..9896bd22ef 100644
--- a/core/partitioning/SegmentedBlock.h
+++ b/core/partitioning/SegmentedBlock.h
@@ -2,9 +2,11 @@
 
 #include <vector>
 
-#include "core/conversion/conversion.h"
+#include "NvInfer.h"
 #include "torch/csrc/jit/ir/ir.h"
 
+#include "core/partitioning/PartitionInfo.h"
+
 namespace trtorch {
 namespace core {
 namespace partitioning {
@@ -111,7 +113,7 @@ struct SegmentedBlock {
 
  private:
   SegmentedBlockTarget target_;
-  std::vector<nvinfer1::Dims> in_shape_;
+  std::vector<nvinfer1::Dims> in_shape_; // REVIEW: This should just be ir::InputRange
   std::vector<torch::jit::Value*> inputs_;
   std::vector<torch::jit::Value*> outputs_;
   std::vector<torch::jit::Node*> nodes_;
@@ -120,10 +122,6 @@ struct SegmentedBlock {
   std::unordered_map<torch::jit::Value*, torch::jit::Value*> old_to_new_;
 };
 
-std::vector<SegmentedBlock> segment_graph(
-    std::shared_ptr<torch::jit::Graph> g,
-    const conversion::TorchFallback& fallback_info);
-
-} // namespace partitioning
+} // namespace ir
 } // namespace core
 } // namespace trtorch
\ No newline at end of file
diff --git a/core/partitioning/partitioning.cpp b/core/partitioning/partitioning.cpp
index 3de69c2464..dde3c674ef 100644
--- a/core/partitioning/partitioning.cpp
+++ b/core/partitioning/partitioning.cpp
@@ -1,7 +1,8 @@
 #include "partitioning.h"
 
 #include <queue>
-#include "shape_analysis.h"
+#include "core/conversion/conversion.h"
+#include "core/partitioning/shape_analysis.h"
 #include "torch/csrc/jit/passes/constant_pooling.h"
 
 namespace trtorch {
@@ -55,7 +56,7 @@ SegmentedBlock injectNodesForNonTensorInputs(SegmentedBlock& seg_block) {
   return std::move(SegmentedBlock(seg_block.target(), new_block_nodes));
 }
 
-void resolveNonTensorInputs(std::vector<SegmentedBlock>& segmented_blocks, std::shared_ptr<torch::jit::Graph> g) {
+void resolveNonTensorInputs(PartitionedGraph& segmented_blocks, std::shared_ptr<torch::jit::Graph> g) {
   // for NonTensor inputs in TensorRT segments, count the usages on Torch segments and TensorRT segments
   std::unordered_map<torch::jit::Value*, usage_info> usage_counts;
   for (int i = segmented_blocks.size() - 1; i >= 0; --i) {
@@ -97,7 +98,7 @@ void resolveNonTensorInputs(std::vector<SegmentedBlock>& segmented_blocks, std::
   return;
 }
 
-void registerSegmentsOutputs(std::vector<SegmentedBlock>& segmented_blocks, std::shared_ptr<torch::jit::Graph> g) {
+void registerSegmentsOutputs(PartitionedGraph& segmented_blocks, std::shared_ptr<torch::jit::Graph> g) {
   // find the corresponding raw values in original global graph for this segmented block's inputs/outputs
   std::set<torch::jit::Value*> input_values;
   for (auto& seg_block : segmented_blocks) {
@@ -152,12 +153,59 @@ void registerSegmentsOutputs(std::vector<SegmentedBlock>& segmented_blocks, std:
   return;
 }
 
+std::vector<SegmentedBlock> segment_graph(
+    std::shared_ptr<torch::jit::Graph> g,
+    const PartitionInfo& partition_info) {
+  auto min_block_size = partition_info.min_block_size;
+  std::unordered_set<std::string> forced_fallback_operators(
+      partition_info.forced_fallback_operators.begin(), partition_info.forced_fallback_operators.end());
+
+  auto nodes = g->block()->nodes();
+  std::vector<SegmentedBlock> segmented_blocks;
+
+  // segment the nodes
+  std::vector<torch::jit::Node*> tensorrt_nodes, pytorch_nodes;
+  for (const auto n : nodes) {
+    if (n->kind() == torch::jit::prim::Constant)
+      continue;
+
+    std::string node_string(n->kind().toQualString());
+    if (conversion::OpSupported(n) && !forced_fallback_operators.count(node_string)) {
+      tensorrt_nodes.push_back(n);
+      if (tensorrt_nodes.size() >= min_block_size && !pytorch_nodes.empty()) {
+        segmented_blocks.emplace_back(SegmentedBlock::kTorch, pytorch_nodes);
+        pytorch_nodes.clear();
+      }
+    } else {
+      if (tensorrt_nodes.size() >= min_block_size) {
+        segmented_blocks.emplace_back(SegmentedBlock::kTensorRT, tensorrt_nodes);
+      } else {
+        pytorch_nodes.insert(pytorch_nodes.end(), tensorrt_nodes.begin(), tensorrt_nodes.end());
+      }
+      tensorrt_nodes.clear();
+      pytorch_nodes.push_back(n);
+    }
+  }
+
+  // if there is any kTorch nodes left, then either the last nodes are kTorch or last nodes are kTensorRT but num <
+  // min_block_size
+  if (!pytorch_nodes.empty()) {
+    pytorch_nodes.insert(pytorch_nodes.end(), tensorrt_nodes.begin(), tensorrt_nodes.end());
+    segmented_blocks.emplace_back(SegmentedBlock::kTorch, pytorch_nodes);
+  } else {
+    segmented_blocks.emplace_back(SegmentedBlock::kTensorRT, tensorrt_nodes);
+  }
+
+  return std::move(segmented_blocks);
+}
+
+
 std::vector<SegmentedBlock> Partition(
     std::shared_ptr<torch::jit::Graph> g,
-    std::vector<conversion::InputRange>& input_ranges,
-    const conversion::TorchFallback& fallback_info) {
+    std::vector<ir::InputRange>& input_ranges,
+    const PartitionInfo& partition_info) {
   // segment lowering global graph into blocks
-  std::vector<SegmentedBlock> segmented_blocks = segment_graph(g, fallback_info);
+  std::vector<SegmentedBlock> segmented_blocks = segment_graph(g, partition_info);
 
   // resolve nonTensor inputs/outputs
   resolveNonTensorInputs(segmented_blocks, g);
diff --git a/core/partitioning/partitioning.h b/core/partitioning/partitioning.h
index fe43240ada..f4ef31785f 100644
--- a/core/partitioning/partitioning.h
+++ b/core/partitioning/partitioning.h
@@ -2,8 +2,8 @@
 
 #include <vector>
 
-#include "core/conversion/conversion.h"
-#include "core/conversion/evaluators/eval_util.h"
+#include "core/ir/ir.h"
+#include "core/partitioning/PartitionInfo.h"
 #include "core/partitioning/SegmentedBlock.h"
 #include "core/util/prelude.h"
 #include "torch/csrc/jit/ir/ir.h"
@@ -12,10 +12,12 @@ namespace trtorch {
 namespace core {
 namespace partitioning {
 
+typedef std::vector<SegmentedBlock> PartitionedGraph;
+
 std::vector<SegmentedBlock> Partition(
     std::shared_ptr<torch::jit::Graph> g,
-    std::vector<conversion::InputRange>& input_ranges,
-    const conversion::TorchFallback& fallback_info);
+    std::vector<ir::InputRange>& input_ranges,
+    const PartitionInfo& partition_info);
 
 } // namespace partitioning
 } // namespace core
diff --git a/core/partitioning/shape_analysis.cpp b/core/partitioning/shape_analysis.cpp
index 513124eb3d..b97613b624 100644
--- a/core/partitioning/shape_analysis.cpp
+++ b/core/partitioning/shape_analysis.cpp
@@ -1,11 +1,12 @@
-#include "shape_analysis.h"
+#include "core/partitioning/shape_analysis.h"
+#include "core/util/prelude.h"
 #include "torch/csrc/jit/api/module.h"
 
 namespace trtorch {
 namespace core {
 namespace partitioning {
 
-std::vector<torch::jit::IValue> generateRandomInputs(std::vector<conversion::InputRange>& input_ranges) {
+std::vector<torch::jit::IValue> generateRandomInputs(std::vector<ir::InputRange>& input_ranges) {
   // generate random inputs for running pytorch segments
   std::vector<torch::jit::IValue> random_inputs;
   for (auto& input_range : input_ranges) {
diff --git a/core/partitioning/shape_analysis.h b/core/partitioning/shape_analysis.h
index 8252573430..0f184cdd62 100644
--- a/core/partitioning/shape_analysis.h
+++ b/core/partitioning/shape_analysis.h
@@ -1,10 +1,12 @@
-#include "SegmentedBlock.h"
+#include "core/partitioning/SegmentedBlock.h"
+#include "core/ir/ir.h"
+
 
 namespace trtorch {
 namespace core {
 namespace partitioning {
 
-std::vector<torch::jit::IValue> generateRandomInputs(std::vector<conversion::InputRange>& input_ranges);
+std::vector<torch::jit::IValue> generateRandomInputs(std::vector<ir::InputRange>& input_ranges);
 
 void getSegmentsOutputByRunning(
     SegmentedBlock& seg_block,
diff --git a/cpp/api/src/compile_spec.cpp b/cpp/api/src/compile_spec.cpp
index dc0c340bb0..3cd047477d 100644
--- a/cpp/api/src/compile_spec.cpp
+++ b/cpp/api/src/compile_spec.cpp
@@ -62,12 +62,12 @@ CompileSpec::CompileSpec(std::vector<std::vector<int64_t>> fixed_sizes) {
   }
 }
 
-core::conversion::InputRange to_internal_input_range(CompileSpec::InputRange i) {
-  return core::conversion::InputRange(i.min, i.opt, i.max);
+core::ir::InputRange to_internal_input_range(CompileSpec::InputRange i) {
+  return core::ir::InputRange(i.min, i.opt, i.max);
 }
 
-std::vector<core::conversion::InputRange> to_vec_internal_input_ranges(std::vector<CompileSpec::InputRange> external) {
-  std::vector<core::conversion::InputRange> internal;
+std::vector<core::ir::InputRange> to_vec_internal_input_ranges(std::vector<CompileSpec::InputRange> external) {
+  std::vector<core::ir::InputRange> internal;
   for (auto range : external) {
     internal.push_back(to_internal_input_range(range));
   }
@@ -96,10 +96,9 @@ core::CompileSpec to_internal_compile_spec(CompileSpec external) {
   internal.convert_info.engine_settings.strict_types = external.strict_types;
   internal.convert_info.engine_settings.device.allow_gpu_fallback = external.device.allow_gpu_fallback;
   internal.convert_info.engine_settings.max_batch_size = external.max_batch_size;
-  internal.convert_info.engine_settings.torch_fallback.enabled = external.torch_fallback.enabled;
-  internal.convert_info.engine_settings.torch_fallback.min_block_size = external.torch_fallback.min_block_size;
-  internal.convert_info.engine_settings.torch_fallback.forced_fallback_operators =
-      external.torch_fallback.forced_fallback_operators;
+  internal.partition_info.enabled = external.torch_fallback.enabled;
+  internal.partition_info.min_block_size = external.torch_fallback.min_block_size;
+  internal.partition_info.forced_fallback_operators = external.torch_fallback.forced_fallback_operators;
 
   switch (external.device.device_type) {
     case CompileSpec::Device::DeviceType::kDLA:
diff --git a/py/trtorch/csrc/tensorrt_classes.cpp b/py/trtorch/csrc/tensorrt_classes.cpp
index 8d76fc1278..2686d3c650 100644
--- a/py/trtorch/csrc/tensorrt_classes.cpp
+++ b/py/trtorch/csrc/tensorrt_classes.cpp
@@ -109,9 +109,9 @@ core::CompileSpec CompileSpec::toInternalCompileSpec() {
   info.convert_info.engine_settings.device.dla_core = device.dla_core;
   info.convert_info.engine_settings.device.allow_gpu_fallback = device.allow_gpu_fallback;
   info.convert_info.engine_settings.torch_fallback.enabled = torch_fallback.enabled;
-  info.convert_info.engine_settings.torch_fallback.min_block_size = torch_fallback.min_block_size;
-  info.convert_info.engine_settings.torch_fallback.forced_fallback_operators = torch_fallback.forced_fallback_operators;
-  info.convert_info.engine_settings.truncate_long_and_double = truncate_long_and_double;
+  info.partition_info.enabled = torch_fallback.enabled;
+  info.partition_info.min_block_size = torch_fallback.min_block_size;
+  info.partition_info.forced_fallback_operators = torch_fallback.forced_fallback_operators;
 
   info.convert_info.engine_settings.capability = toTRTEngineCapability(capability);
   TRTORCH_CHECK(num_min_timing_iters >= 0, "num_min_timing_iters must be 0 or greater");
diff --git a/tests/util/run_graph_engine.cpp b/tests/util/run_graph_engine.cpp
index e39fd80f5f..a12bcff062 100644
--- a/tests/util/run_graph_engine.cpp
+++ b/tests/util/run_graph_engine.cpp
@@ -1,6 +1,7 @@
 #include "NvInfer.h"
 #include "c10/cuda/CUDAStream.h"
 #include "core/conversion/conversion.h"
+#include "core/ir/ir.h"
 #include "core/runtime/runtime.h"
 #include "core/util/prelude.h"
 #include "cuda_runtime_api.h"
@@ -15,16 +16,16 @@ namespace trtorch {
 namespace tests {
 namespace util {
 
-std::vector<core::conversion::InputRange> toInputRanges(std::vector<at::Tensor> ten) {
-  std::vector<core::conversion::InputRange> a;
+std::vector<core::ir::InputRange> toInputRanges(std::vector<at::Tensor> ten) {
+  std::vector<core::ir::InputRange> a;
   for (auto i : ten) {
-    a.push_back(core::conversion::InputRange(core::util::toVec(i.sizes())));
+    a.push_back(core::ir::InputRange(core::util::toVec(i.sizes())));
   }
   return std::move(a);
 }
 
-std::vector<core::conversion::InputRange> toInputRangesDynamic(std::vector<at::Tensor> ten, bool dynamic_batch) {
-  std::vector<core::conversion::InputRange> a;
+std::vector<core::ir::InputRange> toInputRangesDynamic(std::vector<at::Tensor> ten, bool dynamic_batch) {
+  std::vector<core::ir::InputRange> a;
 
   for (auto i : ten) {
     auto opt = core::util::toVec(i.sizes());
@@ -36,7 +37,7 @@ std::vector<core::conversion::InputRange> toInputRangesDynamic(std::vector<at::T
       min_range[0] = ceil(opt[0] / 2.0);
       max_range[0] = 2 * opt[0];
 
-      a.push_back(core::conversion::InputRange(min_range, opt, max_range));
+      a.push_back(core::ir::InputRange(min_range, opt, max_range));
     } else {
       std::vector<int64_t> min_range(opt);
       std::vector<int64_t> max_range(opt);
@@ -44,7 +45,7 @@ std::vector<core::conversion::InputRange> toInputRangesDynamic(std::vector<at::T
       min_range[1] = ceil(opt[1] / 2.0);
       max_range[1] = 2 * opt[1];
 
-      a.push_back(core::conversion::InputRange(min_range, opt, max_range));
+      a.push_back(core::ir::InputRange(min_range, opt, max_range));
     }
   }