diff --git a/Colab-TextGen-GPU.ipynb b/Colab-TextGen-GPU.ipynb
index 3efc9d92..e1950b75 100644
--- a/Colab-TextGen-GPU.ipynb
+++ b/Colab-TextGen-GPU.ipynb
@@ -106,7 +106,7 @@
         "    url_parts = model_url.strip('/').strip().split('/')\n",
         "    output_folder = f\"{url_parts[-2]}_{url_parts[-1]}\"\n",
         "    branch = branch.strip('\"\\' ')\n",
-        "    if branch.strip() != '':\n",
+        "    if branch.strip() != not in ['', 'main']:\n",
         "        output_folder += f\"_{branch}\"\n",
         "        !python download-model.py {model_url} --branch {branch}\n",
         "    else:\n",