Use type parameters

src/lora.jl

@@ -1,7 +1,7 @@
-struct LoRADense
-    primary::Dense
-    proj1::Dense
-    proj2::Dense
+struct LoRADense{P0<:Dense,P1<:Dense,P2<:Dense}
+    primary::P0
+    proj1::P1
+    proj2::P2
 end
 
 """
@@ -10,8 +10,7 @@ end
 Create a LoRA wrapper around a Dense layer. The second projection matrix is initialized to zero, and only the two projections (and not the primary layer) are trainable.
 """
 function LoRADense(primary::Dense, hidden_dim::Int; init=Flux.kaiming_uniform())
-    dim1 = size(primary.weight, 2)
-    dim2 = size(primary.weight, 1)
+    dim2, dim1 = size(primary.weight)
     ld = LoRADense(
         primary,
         Dense(dim1 => hidden_dim, bias=false, init = init),
@@ -21,8 +20,6 @@ function LoRADense(primary::Dense, hidden_dim::Int; init=Flux.kaiming_uniform())
     return ld
 end 
 
-function (lora::LoRADense)(x)
-    return lora.primary(x) .+ lora.proj2(lora.proj1(x))
-end
+(lora::LoRADense)(x) = lora.primary(x) .+ lora.proj2(lora.proj1(x))
 
 Flux.@layer :expand LoRADense trainable=(proj1, proj2)