update notebooks for s3 wheels

Summary: Prepare the tutorial notebooks to use wheels from S3 when run on colab.

Reviewed By: nikhilaravi

Differential Revision: D26226932

fbshipit-source-id: 1f9366c3fb4ba195333a5d5dfa3f6876ea934508

docs/tutorials/fit_simple_neural_radiance_field.ipynb

@@ -38,7 +38,7 @@
    "metadata": {},
    "source": [
     "## 0. Install and Import modules\n",
-    "If `torch` and `pytorch3d` are not installed, run the following cell:"
+    "Ensure `torch` and `torchvision` are installed. If `pytorch3d` is not installed, install it using the following cell:"
    ]
   },
   {
@@ -47,13 +47,29 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# !pip install torch\n",
-    "# import sys\n",
-    "# import torch\n",
-    "# if torch.__version__=='1.6.0+cu101' and sys.platform.startswith('linux'):\n",
-    "#     !pip install pytorch3d\n",
-    "# else:\n",
-    "#     !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
+    "import os\n",
+    "import sys\n",
+    "import torch\n",
+    "need_pytorch3d=False\n",
+    "try:\n",
+    "    import pytorch3d\n",
+    "except ModuleNotFoundError:\n",
+    "    need_pytorch3d=True\n",
+    "if need_pytorch3d:\n",
+    "    if torch.__version__.startswith(\"1.7\") and sys.platform.startswith(\"linux\"):\n",
+    "        # We try to install PyTorch3D via a released wheel.\n",
+    "        version_str=\"\".join([\n",
+    "            f\"py3{sys.version_info.minor}_cu\",\n",
+    "            torch.version.cuda.replace(\".\",\"\"),\n",
+    "            f\"_pyt{torch.__version__[0:5:2]}\"\n",
+    "        ])\n",
+    "        !pip install pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
+    "    else:\n",
+    "        # We try to install PyTorch3D from source.\n",
+    "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
+    "        !tar xzf 1.10.0.tar.gz\n",
+    "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\n",
+    "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
    ]
   },
   {