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 <noreply@anthropic.com>
This commit is contained in:
@@ -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))
|
||||
|
||||
Reference in New Issue
Block a user