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,
         )