Has anyone managed to run evaluate.py with the provided ARC-AGI-2 checkpoint?

train_state.model.load_state_dict(torch.load(eval_cfg.checkpoint, map_location="cuda"), assign=True)

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[81], line 1
----> 1 train_state.model.load_state_dict(torch.load(eval_cfg.checkpoint, map_location="cuda"), assign=True)

File /opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py:2624, in Module.load_state_dict(self, state_dict, strict, assign)
   2616         error_msgs.insert(
   2617             0,
   2618             "Missing key(s) in state_dict: {}. ".format(
   2619                 ", ".join(f'"{k}"' for k in missing_keys)
   2620             ),
   2621         )
   2623 if len(error_msgs) > 0:
-> 2624     raise RuntimeError(
   2625         "Error(s) in loading state_dict for {}:\n\t{}".format(
   2626             self.__class__.__name__, "\n\t".join(error_msgs)
   2627         )
   2628     )
   2629 return _IncompatibleKeys(missing_keys, unexpected_keys)

RuntimeError: Error(s) in loading state_dict for OptimizedModule:
	size mismatch for _orig_mod.model.inner.puzzle_emb.weights: copying a param with shape torch.Size([1045829, 512]) from checkpoint, the shape in current model is torch.Size([1921251, 512]).