vram management fix

2025-03-15 16:40:39 +02:00
parent 3aa573a74f
commit d455cc89a3
2 changed files with 4 additions and 1 deletions
@@ -75,6 +75,9 @@ def enable_vram_management_recursively(model: torch.nn.Module, module_map: dict,
    for name, module in model.named_children():
        for source_module, target_module in module_map.items():
            if isinstance(module, source_module):
+                if "rope_embedder" in name or "patch_embedding" in name:
+                    continue
+
                num_param = sum(p.numel() for p in module.parameters())
                if max_num_param is not None and total_num_param + num_param > max_num_param:
                    module_config_ = overflow_module_config