From 988aa34f3d84ba077311cbf68eb921ed44825d43 Mon Sep 17 00:00:00 2001
From: Nicolas Patry <patry.nicolas@protonmail.com>
Date: Fri, 31 May 2024 22:56:31 +0000
Subject: [PATCH] Fix non decoding paths.

---
 server/text_generation_server/layers/attention/cuda.py | 1 +
 server/text_generation_server/layers/attention/rocm.py | 1 +
 2 files changed, 2 insertions(+)

diff --git a/server/text_generation_server/layers/attention/cuda.py b/server/text_generation_server/layers/attention/cuda.py
index a564237c..d87d9199 100644
--- a/server/text_generation_server/layers/attention/cuda.py
+++ b/server/text_generation_server/layers/attention/cuda.py
@@ -159,6 +159,7 @@ def paged_attention(
                 "auto",
                 1.0,
             )
+    return out
 
 
 try:
diff --git a/server/text_generation_server/layers/attention/rocm.py b/server/text_generation_server/layers/attention/rocm.py
index c5c485de..36db12d0 100644
--- a/server/text_generation_server/layers/attention/rocm.py
+++ b/server/text_generation_server/layers/attention/rocm.py
@@ -129,6 +129,7 @@ def paged_attention(
             "auto",
             1.0,
         )
+    return out
 
 
 if ENGINE != "triton":