BUG: Increasing max_peels parameter causes kilosort4 to crash

[mikemanookin]

Describe the issue:

I have been experimenting with the recently added max_peels parameter in kilosort version 4 (4.0.24) to try to reduce the number of contaminated clusters. Using the default value of 100 works fine, but increasing this value (e.g., 1000 or 10000) causes kilosort to crash. The kilosort log file is attached, but this is the most relevant part of the stack trace:

kilosort4.log

2025-02-03 17:18:26,879 kilosort.template_matching ERROR    Error in template_matching.extract on batch 274
Traceback (most recent call last):
  File "/home/mike/anaconda3/envs/kilosort/lib/python3.9/site-packages/kilosort/template_matching.py", line 84, in extract
    stt, amps, th_amps, Xres = run_matching(ops, X, U, ctc, device=device)
  File "/home/mike/anaconda3/envs/kilosort/lib/python3.9/site-packages/kilosort/template_matching.py", line 222, in run_matching
    st[k:k+nsp, 0] = iX[:,0]
RuntimeError: The expanded size of the tensor (69) must match the existing size (348) at non-singleton dimension 0.  Target sizes: [69].  Tensor sizes: [348]
2025-02-03 17:18:26,879 kilosort.template_matching DEBUG    X shape: torch.Size([519, 60122])
2025-02-03 17:18:26,879 kilosort.template_matching DEBUG    stt shape: torch.Size([60446, 2])
2025-02-03 17:18:26,879 kilosort.run_kilosort ERROR    Encountered error in `run_kilosort`:
Traceback (most recent call last):
  File "/home/mike/anaconda3/envs/kilosort/lib/python3.9/site-packages/kilosort/run_kilosort.py", line 242, in run_kilosort
    st,tF, _, _ = detect_spikes(
  File "/home/mike/anaconda3/envs/kilosort/lib/python3.9/site-packages/kilosort/run_kilosort.py", line 673, in detect_spikes
    st, tF, ops = template_matching.extract(ops, bfile, Wall3, device=device,
  File "/home/mike/anaconda3/envs/kilosort/lib/python3.9/site-packages/kilosort/template_matching.py", line 84, in extract
    stt, amps, th_amps, Xres = run_matching(ops, X, U, ctc, device=device)
  File "/home/mike/anaconda3/envs/kilosort/lib/python3.9/site-packages/kilosort/template_matching.py", line 222, in run_matching
    st[k:k+nsp, 0] = iX[:,0]
RuntimeError: The expanded size of the tensor (69) must match the existing size (348) at non-singleton dimension 0.  Target sizes: [69].  Tensor sizes: [348]

Reproduce the bug:

Error message:

Version information:

Kilosort version 4.0.24