fix: fix triton OutOfResources import

773aabdd · OlivierDehaene · 50b495f3 · 773aabdd
Commit 773aabdd authored Dec 14, 2023 by OlivierDehaene
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

server/text_generation_server/utils/gptq/custom_autotune.py server/text_generation_server/utils/gptq/custom_autotune.py +1 -1

No files found.
--- 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):