fix device error

This commit is contained in:
Kohya S
2023-03-11 08:03:02 +09:00
parent b80431de30
commit 8b25929765
2 changed files with 9 additions and 0 deletions

View File

@@ -77,6 +77,12 @@ def merge_lora_models(models, ratios, new_rank, new_conv_rank, device, merge_dty
# W <- W + U * D
scale = (alpha / network_dim)
if device: # and isinstance(scale, torch.Tensor):
scale = scale.to(device)
up_weight = up_weight.to(device)
down_weight = down_weight.to(device)
if not conv2d: # linear
weight = weight + ratio * (up_weight @ down_weight) * scale
elif kernel_size == (1, 1):