fix cuda issues

Author: sokrypton

rf/examples/diffusion.ipynb

@@ -38,7 +38,7 @@
       },
       "outputs": [],
       "source": [
-        "#@title setup **RFdiffusion** (~2m)\n",
+        "#@title setup **RFdiffusion** (~3min)\n",
         "%%time\n",
         "import os, time, signal\n",
         "import sys, random, string, re\n",
@@ -58,9 +58,12 @@
         "if not os.path.isdir(\"RFdiffusion\"):\n",
         "  print(\"installing RFdiffusion...\")\n",
         "  os.system(\"git clone https://github.com/sokrypton/RFdiffusion.git\")\n",
-        "  os.system(\"pip -q install jedi omegaconf hydra-core icecream pyrsistent\")\n",
-        "  os.system(\"pip install dgl -f https://data.dgl.ai/wheels/cu121/repo.html\")\n",
-        "  os.system(\"cd RFdiffusion/env/SE3Transformer; pip -q install --no-cache-dir -r requirements.txt; pip -q install .\")\n",
+        "  os.system(\"pip install jedi omegaconf hydra-core icecream pyrsistent pynvml decorator\")\n",
+        "  os.system(\"pip install git+https://github.com/NVIDIA/dllogger#egg=dllogger\")\n",
+        "  # 17Mar2024: adding --no-dependencies to avoid installing nvidia-cuda-* dependencies\n",
+        "  os.system(\"pip install --no-dependencies dgl==2.0.0 -f https://data.dgl.ai/wheels/cu121/repo.html\")\n",
+        "  os.system(\"pip install --no-dependencies e3nn==0.3.3 opt_einsum_fx\")\n",
+        "  os.system(\"cd RFdiffusion/env/SE3Transformer; pip install .\")\n",
         "  os.system(\"wget -qnc https://files.ipd.uw.edu/krypton/ananas\")\n",
         "  os.system(\"chmod +x ananas\")\n",
         "\n",