From 501639b099b5ebcdabc3ada329faabff91a52561 Mon Sep 17 00:00:00 2001 From: Aniket Maurya Date: Fri, 10 Jan 2025 12:44:20 +0000 Subject: [PATCH] fix callback runner to execute after predict (#406) --- src/litserve/loops/simple_loops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/litserve/loops/simple_loops.py b/src/litserve/loops/simple_loops.py index cd5014de..538d7d15 100644 --- 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)