fix callback runner to execute after predict (#406)

Lightning-AI · Jan 10, 2025 · 501639b · 501639b
1 parent 747308a
commit 501639b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/litserve/loops/simple_loops.py b/src/litserve/loops/simple_loops.py
@@ -197,7 +197,7 @@ def run_batched_loop(
 
             callback_runner.trigger_event(EventTypes.BEFORE_PREDICT, lit_api=lit_api)
             y = _inject_context(contexts, lit_api.predict, x)
-            callback_runner.trigger_event(EventTypes.BEFORE_PREDICT, lit_api=lit_api)
+            callback_runner.trigger_event(EventTypes.AFTER_PREDICT, lit_api=lit_api)
 
             outputs = lit_api.unbatch(y)