diff --git a/notebooks/Interacting_with_CLIP.ipynb b/notebooks/Interacting_with_CLIP.ipynb
index b13a90e..4975e65 100644
--- a/notebooks/Interacting_with_CLIP.ipynb
+++ b/notebooks/Interacting_with_CLIP.ipynb
@@ -352,10 +352,11 @@
       "source": [
         "import numpy as np\n",
         "import torch\n",
+        "from pkg_resources import packaging\n",
         "\n",
         "print(\"Torch version:\", torch.__version__)\n",
         "\n",
-        "assert torch.__version__.split(\".\") >= [\"1\", \"7\", \"1\"], \"PyTorch 1.7.1 or later is required\""
+        "assert packaging.version.parse(torch.__version__) >= packaging.version.parse(\"1.7.1\"), \"PyTorch 1.7.1 or later is required\""
       ],
       "execution_count": 2,
       "outputs": [
diff --git a/notebooks/Prompt_Engineering_for_ImageNet.ipynb b/notebooks/Prompt_Engineering_for_ImageNet.ipynb
index 1dde052..5b5ca92 100644
--- a/notebooks/Prompt_Engineering_for_ImageNet.ipynb
+++ b/notebooks/Prompt_Engineering_for_ImageNet.ipynb
@@ -588,10 +588,11 @@
         "import torch\n",
         "import clip\n",
         "from tqdm.notebook import tqdm\n",
+        "from pkg_resources import packaging\n",
         "\n",
         "print(\"Torch version:\", torch.__version__)\n",
         "\n",
-        "assert torch.__version__.split(\".\") >= [\"1\", \"7\", \"1\"], \"PyTorch 1.7.1 or later is required\""
+        "assert packaging.version.parse(torch.__version__) >= packaging.version.parse(\"1.7.1\"), \"PyTorch 1.7.1 or later is required\""
       ],
       "execution_count": 2,
       "outputs": [
@@ -1105,4 +1106,4 @@
       ]
     }
   ]
-}
\ No newline at end of file
+}