From a5678c91858076ee1e3e7d82ee62a748d6541b63 Mon Sep 17 00:00:00 2001 From: ccx1324 <202431060119@qq.com> Date: Mon, 20 Jul 2026 16:11:48 +0800 Subject: [PATCH] fix: create LoRA parameters on base weight device instead of CPU When `inject_lora()` replaces Linear layers with LoRALinear after the model has been moved to CUDA, the new lora_A and lora_B parameters were always created on CPU, causing a device mismatch error during the forward pass. Now lora_A and lora_B are created on the same device and dtype as the parent weight, matching the model's current device. Co-Authored-By: Claude Opus 4.7 --- astrai/model/components/lora.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/astrai/model/components/lora.py b/astrai/model/components/lora.py index 1850663..f7cc4a8 100644 --- a/astrai/model/components/lora.py +++ b/astrai/model/components/lora.py @@ -39,8 +39,10 @@ class LoRALinear(nn.Module): self.r = r self.scaling = alpha / r - self.lora_A = nn.Parameter(torch.randn(r, self.weight.shape[1]) / r) - self.lora_B = nn.Parameter(torch.zeros(self.weight.shape[0], r)) + device = self.weight.device + dtype = self.weight.dtype + self.lora_A = nn.Parameter(torch.randn(r, self.weight.shape[1], device=device, dtype=dtype) / r) + self.lora_B = nn.Parameter(torch.zeros(self.weight.shape[0], r, device=device, dtype=dtype)) self._merged = False def forward(self, x):