From 773aabdda6197cae3d2092f1cd6d9ce08d649185 Mon Sep 17 00:00:00 2001
From: OlivierDehaene <23298448+OlivierDehaene@users.noreply.github.com>
Date: Thu, 14 Dec 2023 16:04:26 +0100
Subject: [PATCH] fix: fix triton OutOfResources import

---
 server/text_generation_server/utils/gptq/custom_autotune.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/server/text_generation_server/utils/gptq/custom_autotune.py b/server/text_generation_server/utils/gptq/custom_autotune.py
index 17dff02e..589d89ef 100644
--- a/server/text_generation_server/utils/gptq/custom_autotune.py
+++ b/server/text_generation_server/utils/gptq/custom_autotune.py
@@ -90,7 +90,7 @@ class Autotuner(triton.KernelInterface):
             return triton.testing.do_bench(
                 kernel_call, percentiles=(0.5, 0.2, 0.8), rep=40
             )
-        except triton.compiler.OutOfResources:
+        except triton.OutOfResources:
             return (float("inf"), float("inf"), float("inf"))
 
     def run(self, *args, **kwargs):