i am trying to use model for inference but getting this error while running the script python script is same as colab demo code can someone help me with this
edit
i am using jax[tpu]>=0.2.16 cause when i tried using 0.2.12 as given in colab i am getting error saying
i am trying to use model for inference but getting this error while running the script python script is same as colab demo code can someone help me with this
edit i am using jax[tpu]>=0.2.16 cause when i tried using 0.2.12 as given in colab i am getting error saying![Screenshot from 2022-02-08 21-37-30](https://user-images.githubusercontent.com/72891217/153027354-fbfa3685-5692-4314-a968-f018025ae5e7.png)