add setter of max_new_tokens

deepspeedai · Jan 29, 2024 · 467d768 · 467d768
1 parent 7bc17c2
commit 467d768
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/mii/batching/data_classes.py b/mii/batching/data_classes.py
@@ -88,6 +88,10 @@ def min_new_tokens(self) -> int:
     def max_new_tokens(self) -> int:
         return self.generate_params.max_new_tokens
 
+    @max_new_tokens.setter
+    def max_new_tokens(self, max_new_tokens: int) -> None:
+        self.generate_params.max_new_tokens = max_new_tokens
+
     @property
     def stream(self) -> bool:
         return self.generate_params.stream