attach batch_size

Lightning-AI · Dec 10, 2024 · a822eac · a822eac
1 parent 7aff112
commit a822eac
Showing 1 changed file with 1 addition and 0 deletions.
diff --git 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: