in Google Colab Pro+, flax.errors.ScopeParamShapeError: Inconsistent shap es between value and initializer for parameter "embedding" in "/model/encoder/em bed_tokens"

[dmd]

I tried to run the notebook in Colab Pro+, and got:

UnfilteredStackTrace                      Traceback (most recent call last)
[<ipython-input-17-4a86f506c701>](https://localhost:8080/#) in <module>()
     19         temperature,
---> 20         cond_scale,
     21     )

46 frames
UnfilteredStackTrace: flax.errors.ScopeParamShapeError: Inconsistent shapes between value and initializer for parameter "embedding" in "/model/encoder/embed_tokens": (1, 50272, 2048), (50272, 2048). (https://flax.readthedocs.io/en/latest/flax.errors.html#flax.errors.ScopeParamShapeError)

[borisdayma]

That is strange it seems to be working for me

[technillogue]

I also have this exact issue (similarly to #214, #245), but only running on gpu; bafflingly, it's fine on cpu.

jax                          0.3.13
jaxlib                       0.3.10+cuda11.cudnn82

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.129.06   Driver Version: 470.129.06   CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  Off  | 00000000:01:00.0 Off |                  N/A |
| N/A   49C    P0    25W /  N/A |      5MiB / 16125MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+