diff --git a/src/enzyme_ad/jax/Passes/AlwaysInliner.h b/src/enzyme_ad/jax/Passes/AlwaysInliner.h
index 3f92a80916..169b164c7d 100644
--- a/src/enzyme_ad/jax/Passes/AlwaysInliner.h
+++ b/src/enzyme_ad/jax/Passes/AlwaysInliner.h
@@ -48,6 +48,12 @@ struct AlwaysInlinerInterface : public mlir::InlinerInterface {
     return true;
   }
 
+  bool allowSingleBlockOptimization(
+      iterator_range<Region::iterator> inlinedBlocks) const final {
+    return !(!inlinedBlocks.empty() &&
+             isa<LLVM::UnreachableOp>(inlinedBlocks.begin()->getTerminator()));
+  }
+
   //===--------------------------------------------------------------------===//
   // Transformation Hooks
   //===--------------------------------------------------------------------===//