Fixing in python 3.8

erfanzar · Oct 4, 2023 · d5f2cad · d5f2cad
1 parent f9c59cd
commit d5f2cad
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/lib/python/EasyDel/modules/mistral/modelling_mistral_flax.py b/lib/python/EasyDel/modules/mistral/modelling_mistral_flax.py
@@ -1,4 +1,5 @@
 import functools
+import typing
 
 import flax.core
 from jax import jit, random, grad, numpy as jnp, Array
@@ -648,7 +649,7 @@ def __call__(
             output_hidden_states: bool = False,
             return_dict: bool = True,
 
-    ) -> Tuple[Array, ...] | FlaxBaseModelOutput:
+    ) -> typing.Union[Tuple[Array, ...], FlaxBaseModelOutput]:
         if input_embeds is None:
             input_embeds = self.embed_tokens(input_ids.astype("i4"))