From 8fade1b0e33763dc12bafa6934547d4a529d9db5 Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Sun, 5 Apr 2026 22:17:29 +0200 Subject: [PATCH] fix: initialize LoRA params on same device as wrapped linear apply_lora() is called after generator.to(device), so lora_A/lora_B were being created on CPU while the rest of the model was on CUDA. Co-Authored-By: Claude Sonnet 4.6 --- selva_core/model/lora.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/selva_core/model/lora.py b/selva_core/model/lora.py index c3b0c29..f726cdb 100644 --- a/selva_core/model/lora.py +++ b/selva_core/model/lora.py @@ -42,8 +42,9 @@ class LoRALinear(nn.Module): linear.bias.requires_grad_(False) ref_dtype = linear.weight.dtype - self.lora_A = nn.Parameter(torch.empty(rank, in_f, dtype=ref_dtype)) - self.lora_B = nn.Parameter(torch.zeros(out_f, rank, dtype=ref_dtype)) + ref_device = linear.weight.device + self.lora_A = nn.Parameter(torch.empty(rank, in_f, dtype=ref_dtype, device=ref_device)) + self.lora_B = nn.Parameter(torch.zeros(out_f, rank, dtype=ref_dtype, device=ref_device)) self.scale = alpha / rank nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))