fix bug in subnetwork extract test

whittle-org · Jun 25, 2024 · 70d7de1 · 70d7de1
1 parent f2828b0
commit 70d7de1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/test/test_extract.py b/test/test_extract.py
@@ -13,7 +13,7 @@ def test_extract_sub_network() -> None:
 
     super_network = GPT(config)
     sub_network_config = Config.from_name("pythia-14m")
-
+    sub_network_config.fix_head_size = False
     super_network.eval()
     super_network.set_sub_network(
         sub_network_n_embd=sub_network_config.n_embd,