Fix fused models for tf >= 4.39 (#418)

Author: TechxGenus

awq/modules/fused/model.py

@@ -83,6 +83,14 @@ def __init__(self, vocab_size, blocks, embedding, norm):
         self.blocks: List[LlamaLikeBlock] = nn.ModuleList(blocks)
         self.norm = norm
         self.last_forward_num_tokens = 0
+        
+    @property
+    def embed_tokens(self):
+        return self.embedding
+    
+    @property
+    def layers(self):
+        return self.blocks
 
     @torch.inference_mode()
     def forward(

setup.py

@@ -86,7 +86,7 @@ def get_kernels_whl_url(
 
 requirements = [
     "torch>=2.0.1",
-    "transformers>=4.35.0,<=4.38.2",
+    "transformers>=4.35.0",
     "tokenizers>=0.12.1",
     "typing_extensions>=4.8.0",
     "accelerate",