From 8ef09f8c5c2e90e407ea3302fd9384a64c60eb6f Mon Sep 17 00:00:00 2001
From: Julio <jperez@nvidia.com>
Date: Tue, 29 Nov 2022 10:53:27 -0500
Subject: [PATCH] update triton to allocate more cuda memory to pool size

---
 merlin/systems/triton/utils.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/merlin/systems/triton/utils.py b/merlin/systems/triton/utils.py
index 76127c55a..6d5c1ceb2 100644
--- a/merlin/systems/triton/utils.py
+++ b/merlin/systems/triton/utils.py
@@ -66,6 +66,7 @@ def run_triton_server(
         f"--backend-config={backend_config}",
         f"--grpc-port={grpc_port}",
         f"--grpc-address={grpc_host}",
+        "--cuda-memory-pool-byte-size=0:536870912",
     ]
     env = os.environ.copy()
     env["CUDA_VISIBLE_DEVICES"] = "0"