KeyError in entropix.ipynb

[koolkao]

Appreciate the Colab notebook on Entropix, can't wait to try it out myself! When I run it in Colab, I'm getting this KeyError below.

[<ipython-input-6-d0ff32815be4>](https://localhost:8080/#) in load_weights(ckpt_dir, n_layers, debug)
     54         for i in range(n_layers):
     55             layer_weights.append(LayerWeights(
---> 56                 wq=w[f'layers.{i}.attention.wq.weight'],
     57                 wk=w[f'layers.{i}.attention.wk.weight'],
     58                 wv=w[f'layers.{i}.attention.wv.weight'],

KeyError: 'layers.0.attention.wq.weight'

[SinatrasC]

Make sure you removed comment # from model download call

#download_weights() -> download_weights()

[koolkao]

thank you!