Merge branch 'vchen/fix_clip_trtllm' of https://github.com/xuanzic/NeMo…

… into vchen/fix_clip_trtllm
xuanzic · Aug 7, 2024 · 81b6584 · 81b6584
2 parents 18da971 + 1773287
commit 81b6584
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/nemo/export/multimodal/build.py b/nemo/export/multimodal/build.py
@@ -208,7 +208,10 @@ def forward(self, images):
             return vision_x
 
     encoder = AutoModel.from_pretrained(
-        vision_config["from_pretrained"], torch_dtype=torch.bfloat16, trust_remote_code=True, attn_implementation='eager'
+        vision_config["from_pretrained"],
+        torch_dtype=torch.bfloat16,
+        trust_remote_code=True,
+        attn_implementation='eager',
     )
     vision_encoder = encoder.vision_model
     hf_config = encoder.config
@@ -326,7 +329,10 @@ def forward(self, images):
             return vision_x
 
     encoder = AutoModel.from_pretrained(
-        vision_config["from_pretrained"], torch_dtype=torch.bfloat16, trust_remote_code=True, attn_implementation='eager'
+        vision_config["from_pretrained"],
+        torch_dtype=torch.bfloat16,
+        trust_remote_code=True,
+        attn_implementation='eager',
     )
     vision_encoder = encoder.vision_model
     hf_config = encoder.config