Method apply_max_norm

lycoris/modules/locon.py:274–284

(self, max_norm, device=None)

            if self.module_type == "linear":
                mid = F.linear(x, self.lora_down.weight.to(x))
            elif self.tucker:
                mid = self.lora_mid(self.lora_down(x))
            else:
                mid = self.lora_down(x)
            drop = (
                torch.rand(self.lora_dim, device=mid.device) > self.rank_dropout
            ).to(mid.dtype)
            if self.rank_dropout_scale:
                drop /= drop.mean()
            if (dims := len(x.shape)) == 4:
                drop = drop.view(1, -1, 1, 1)
            else:
                drop = drop.view(*[1] * (dims - 1), -1)
            mid = mid * drop