diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py
index cbde956a88..8ff1fa9959 100644
--- a/py/torch_tensorrt/dynamo/_compiler.py
+++ b/py/torch_tensorrt/dynamo/_compiler.py
@@ -7,6 +7,7 @@
 import warnings
 from typing import Any, Collection, List, Optional, Sequence, Set, Tuple, Union
 
+import psutil
 import torch
 from torch.export import ExportedProgram
 from torch.fx.node import Target
@@ -694,6 +695,13 @@ def compile(
             else DiskEngineCache(engine_cache_dir, engine_cache_size)
         )
 
+    if cpu_memory_budget and cpu_memory_budget > psutil.virtual_memory().available:
+        logger.warning(
+            f"CPU memory budget is greater than the available memory: {cpu_memory_budget} > {psutil.virtual_memory().available}"
+            "\nUsing the available memory instead."
+        )
+        cpu_memory_budget = psutil.virtual_memory().available
+
     compilation_options = {
         "enabled_precisions": (
             enabled_precisions if enabled_precisions else _defaults.ENABLED_PRECISIONS
diff --git a/tests/py/dynamo/partitioning/test_000_resource_partitioning.py b/tests/py/dynamo/partitioning/test_000_resource_partitioning.py
index 2014eea8fe..b50c5f45d4 100644
--- a/tests/py/dynamo/partitioning/test_000_resource_partitioning.py
+++ b/tests/py/dynamo/partitioning/test_000_resource_partitioning.py
@@ -86,7 +86,7 @@ def forward(self, x):
             current_subgraph = []
             # Split the subgraph into two subgraphs by the ReLU node, which breaks the fusion group.
             for node in subgraphs[0].nodes:
-                if node.op == "call_function" and node.target == aten.relu.default:
+                if node.op == "call_function" and "relu" in str(node.target):
                     new_subgraphs.append(Subgraph(is_acc=True, nodes=current_subgraph))
                     current_subgraph = []
                 current_subgraph.append(node)