added the pytorch cuda cell to the torch_to_jax demo

vaatsalya123 · Feb 20, 2024 · d5dd40c · d5dd40c
1 parent a806ae3
commit d5dd40c
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/examples_and_demos/torch_to_jax.ipynb b/examples_and_demos/torch_to_jax.ipynb
@@ -188,6 +188,19 @@
         "After transpiling our model, we can see what's the improvement in runtime efficiency like. For this let's compile the original PyTorch model using `torch.compile`"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ref : https://github.com/pytorch/pytorch/issues/107960\n",
+        "!export LC_ALL=\"en_US.UTF-8\"\n",
+        "!export LD_LIBRARY_PATH=\"/usr/lib64-nvidia\"\n",
+        "!export LIBRARY_PATH=\"/usr/local/cuda/lib64/stubs\"\n",
+        "!ldconfig /usr/lib64-nvidia"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": 6,