From d5dd40c78269c17561f1a489e2186b97738d1cdf Mon Sep 17 00:00:00 2001
From: Ved Patwardhan <54766411+vedpatwardhan@users.noreply.github.com>
Date: Tue, 20 Feb 2024 11:53:48 +0300
Subject: [PATCH] added the pytorch cuda cell to the torch_to_jax demo

---
 examples_and_demos/torch_to_jax.ipynb | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/examples_and_demos/torch_to_jax.ipynb b/examples_and_demos/torch_to_jax.ipynb
index 63c142af..0d822773 100644
--- 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,