Fix qwen-vl-chat hung (#1824)

InternLM · Jun 21, 2024 · 4067cb2 · 4067cb2
1 parent 21533e7
commit 4067cb2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/lmdeploy/vl/model/qwen.py b/lmdeploy/vl/model/qwen.py
@@ -37,10 +37,10 @@ def build_model(self):
         max_memory = get_balanced_memory(
             model,
             dtype=torch.half,
-            no_split_module_classes=['VisualAttentionBlock'])
+            no_split_module_classes=['VisualAttentionBlock', 'Resampler'])
         device_map = infer_auto_device_map(
             model,
-            no_split_module_classes=['VisualAttentionBlock'],
+            no_split_module_classes=['VisualAttentionBlock', 'Resampler'],
             max_memory=max_memory,
             dtype=torch.half)
         same_device_keys = [('transformer.visual.conv1',