Update vllm/model_executor/model_loader/loader.py

Co-authored-by: Simon Mo <simon.mo@hey.com>
Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>

Author: LucasWilkinson

vllm/model_executor/model_loader/loader.py

@@ -644,7 +644,7 @@ def load_model(self, vllm_config: VllmConfig) -> nn.Module:
                         hasattr(module, "process_weights_after_loading"):
                         # When attention modules need to process weights after
                         # currently only used by MLA
-                        module.process_weights_after_loading()
+                        module.process_weights_after_loading(model_config.dtype)
             rank = get_tensor_model_parallel_rank()
             pattern = os.path.join(
                 local_model_path,