tutorial_predict_next_steps incompatible with model checkpoint & fine-tuned model for time step prediction

[matteomattersberger]

Dear authors,

Thank you for your groundbreaking contribution!

I have been following the tutorial tutorial_predict_next_steps, which refers to the checkpoint at:

BrainLM/training-runs/2023-04-21-10_29_26/checkpoint-181

However, this checkpoint does not appear to be available in the repository. The available checkpoint:

BrainLM/pretrained_models/2023-06-06-22_15_00-checkpoint-1400

seems to have a different architecture, making it incompatible with the data preprocessing steps in the tutorial. Specifically, the old checkpoint appears to expect inputs in the format:

(signal_embedding_projection): Linear(in_features=1, out_features=128, bias=True)
(xyzt_embedding_projection): Linear(in_features=4, out_features=128, bias=True)

whereas the new checkpoint requires:

(signal_embedding_projection): Linear(in_features=20, out_features=256, bias=True)
(xyz_embedding_projection): Linear(in_features=3, out_features=256, bias=True)

Would it be possible to share the older checkpoint or update the tutorial to align with the newer version?

Additionally, I could not locate the fine-tuned model for timepoint prediction mentioned in section 4.3 of the paper. Could you share that model as well?

Once again, many thanks for sharing this exceptional work!

Best regards,
Matteo Mattersberger