Upload utils.py

utils.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
 from transformers.models.gpt2.modeling_gpt2 import GPT2Attention, GPT2MLP
-from typing import Optional, Tuple, Union, List
+from typing import Optional, Tuple, Union, Any, Dict, List
 from transformers import GPT2LMHeadModel
 import torch
 from transformers.deepspeed import is_deepspeed_zero3_enabled