From a822eace22c1837a605185944206fa97459bb344 Mon Sep 17 00:00:00 2001
From: Aniket Maurya <aniketmaurya@outlook.com>
Date: Tue, 10 Dec 2024 13:34:02 +0000
Subject: [PATCH] attach batch_size

---
 src/litserve/api.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/litserve/api.py b/src/litserve/api.py
index 321ea9c2..c54a7b80 100644
--- a/src/litserve/api.py
+++ b/src/litserve/api.py
@@ -114,6 +114,7 @@ def device(self, value):
         self._device = value
 
     def _sanitize(self, max_batch_size: int, spec: Optional[LitSpec]):
+        self.max_batch_size = max_batch_size
         if self.stream:
             self._default_unbatch = self._unbatch_stream
         else: