fix colab

inference_pretrained.ipynb

@@ -12,6 +12,25 @@
     "First, let's start with a minimal example!"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ca93d021",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# run this block if you're using Colab\n",
+    "\n",
+    "# Download repo\n",
+    "!git clone https://github.com/rail-berkeley/crossformer.git\n",
+    "%cd crossformer\n",
+    "# Install repo\n",
+    "!pip3 install -e .\n",
+    "!pip3 install -r requirements.txt\n",
+    "!pip3 install --upgrade \"jax[cuda11_pip]==0.4.20\" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
+    "!pip install numpy==1.21.1 # to fix colab AttributeError: module 'numpy' has no attribute '_no_nep50_warning', if the error still shows reload"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,