From ada722c9e45baa188225bb9a3022e307f4a3fcb1 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Sat, 15 Jun 2024 08:34:33 -0600
Subject: [PATCH] Fixed issue with heads not being added

---
 toolkit/models/ilora.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/toolkit/models/ilora.py b/toolkit/models/ilora.py
index 134f2071..eaf3d29c 100644
--- a/toolkit/models/ilora.py
+++ b/toolkit/models/ilora.py
@@ -284,6 +284,7 @@ class InstantLoRAModule(torch.nn.Module):
             hidden_size=head_dim,
             head_size=head_dim,
             num_mlp_layers=1,
+            num_heads=self.num_heads,
             output_size=self.output_size,
         )