Update utils_llama.py

research/long-context-llama/H2O/utils_llama.py

@@ -1,5 +1,5 @@
 import math
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import pdb
 import types