Tune text_generation.py compilation (#1137)

* Tune text_generation.py compilation * Add os
neuralmagic · Jul 21, 2023 · fb9d131 · fb9d131
1 parent 27d9b82
commit fb9d131
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/deepsparse/transformers/pipelines/text_generation.py b/src/deepsparse/transformers/pipelines/text_generation.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import logging
+import os
 import warnings
 from typing import List, Optional, Tuple, Type, Union
 
@@ -149,6 +150,8 @@ def __init__(
                 "The multi-token engine will not be "
                 "used for prompt processing."
             )
+            if "WAND_OPT_FLAGS" not in os.environ:
+                os.environ["WAND_OPT_FLAGS"] = "default,~pyramids"
 
         self.deterministic = deterministic
         self.sampling_temperature = sampling_temperature