From 06fc4c6ca82f3b0c9a894b0bd6ffd4dd8ff06f82 Mon Sep 17 00:00:00 2001 From: henk717 Date: Sat, 16 Apr 2022 00:02:10 +0200 Subject: [PATCH] Add TPU warning --- colab/TPU.ipynb | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/colab/TPU.ipynb b/colab/TPU.ipynb index 916012b2..9efaaa56 100644 --- a/colab/TPU.ipynb +++ b/colab/TPU.ipynb @@ -53,8 +53,17 @@ "Drive = \"Unextracted (Less Space)\" #@param [\"Unextracted (Less Space)\", \"Extracted (Faster Loading)\"]\n", "#@markdown Extracted models take up more space but load faster the next time you use them, not all models use your Google Drive. See the Model list below for descriptions and space requirements. If your extracted model does not load the next time you try to launch KoboldAI delete the folder from your Google Drive and ensure enough space is available.\n", "\n", + "import os\n", + "try:\n", + " device_name = os.environ['COLAB_TPU_ADDR']\n", + " TPU_ADDRESS = 'grpc://' + device_name\n", + " print('Found TPU at: {}'.format(TPU_ADDRESS))\n", + "except KeyError:\n", + " raise RuntimeError(\"⚠️You can not run this notebook without the TPU accelerator, go to Runtime->Sessions, terminate your session and then try again.⚠️\")\n", + "print('Now we will need your Google Drive to store settings and saves, you must login with the same account you used for Colab.')\n", "from google.colab import drive\n", "drive.mount('/content/drive/')\n", + "\n", "!wget https://henk.tech/ckds -O - | bash /dev/stdin -i drive\n", "\n", "if Model == \"Skein 6B\":\n",