Merge pull request #19 from ccx1324/lora-device-fix

fix: create LoRA parameters on base weight device instead of CPU
2026-07-20 16:14:30 +08:00
parent 2c50b3cf37 a5678c9185
commit 121a7bf8b4
1 changed files with 4 additions and 2 deletions
@@ -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):