Update chess-gpt-eval-contrastive/mamba_module.py
Browse files
chess-gpt-eval-contrastive/mamba_module.py
CHANGED
@@ -104,7 +104,7 @@ class MambaPlayer:
|
|
104 |
self.linear_probes = torch.load(linear_probe_path)
|
105 |
else:
|
106 |
self.linear_probes = {}
|
107 |
-
|
108 |
layer_idx: {
|
109 |
probe_type: optim.Adam(self.linear_probes[layer_idx][probe_type].parameters(), lr=lr)
|
110 |
for probe_type in ['q_value', 'q_value_delta', 'material_balance']
|
|
|
104 |
self.linear_probes = torch.load(linear_probe_path)
|
105 |
else:
|
106 |
self.linear_probes = {}
|
107 |
+
self.linear_optimizers = {
|
108 |
layer_idx: {
|
109 |
probe_type: optim.Adam(self.linear_probes[layer_idx][probe_type].parameters(), lr=lr)
|
110 |
for probe_type in ['q_value', 'q_value_delta', 'material_balance']
|