Update base.py

Lightning-AI · Jan 8, 2025 · f2c6672 · f2c6672
1 parent 0642673
commit f2c6672
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/src/litserve/loops/base.py b/src/litserve/loops/base.py
@@ -265,12 +265,11 @@ def pre_setup(self, lit_api: LitAPI, spec: Optional[LitSpec]):
 
         original = lit_api.unbatch.__code__ is LitAPI.unbatch.__code__
         if not lit_api.stream and any([
-            inspect.isgeneratorfunction(lit_api.predict),
-            inspect.isgeneratorfunction(lit_api.encode_response),
-            (original or inspect.isgeneratorfunction(lit_api.unbatch)),
-        ]):
+                inspect.isgeneratorfunction(lit_api.predict),
+                inspect.isgeneratorfunction(lit_api.encode_response),
+            ]):
             raise ValueError(
-                """When `stream=False`, `lit_api.predict`, `lit_api.encode_response` and `lit_api.unbatch` must not be
+                """When `stream=False`, `lit_api.predict`, `lit_api.encode_response` must not be
                 generator functions.
 
                 Correct example: