Update chess-gpt-eval-contrastive/mamba_module.py

chess-gpt-eval-contrastive/mamba_module.py

@@ -126,7 +126,7 @@ class MambaPlayer:
                         tensor_output = output
                     seq_len = tensor_output.shape[1]
                     bucket = next(b for b in self.move_buckets if self.move_num <= b)
-                    self.activations_sum[layer_idx][bucket]["current"][:, :8, :] += tensor_output.detach().cpu().numpy()[:, :max(self.seq_len, 8), :][:, -8:, :]
+                    self.activations_sum[layer_idx][bucket]["current"][:, :min(8, self.seq_len), :] += tensor_output.detach().cpu().numpy()[:, :self.seq_len, :][:, -8:, :]
                     self.activations_count[layer_idx][bucket]["current"] += 1
                 
                 self.hooks.append(layer.register_forward_hook(hook))
@@ -325,8 +325,8 @@ class MambaPlayer:
         def get_lr(it):
             warmup_iters = 0 #300 * 43
             lr_decay_iters = 3000 * 43
-            learning_rate = 0.0002
-            min_lr = 0.00001
+            learning_rate = 0.000075
+            min_lr = 0.0000075
             # 1) linear warmup for warmup_iters steps
             if it < warmup_iters:
                 return learning_rate * it / warmup_iters