dtype dropdown in colab

2022-07-10 13:23:42 -04:00
parent 1ffdef9a56
commit cfb9f60b6e
2 changed files with 4 additions and 3 deletions
@@ -102,7 +102,7 @@
      },
      "source": [
        "### Load Model\n",
-        "Float32 is faster but uses more GPU memory.  Change the `grid_size` to 3 or less if using float32."
+        "`float32` is faster than `float16` but uses more GPU memory.  Change the `grid_size` to 3 or less if using `float32`."
      ]
    },
    {
@@ -128,13 +128,14 @@
        }
      ],
      "source": [
+        "dtype = \"float32\" #@param [\"float32\", \"float16\", \"bfloat16\"]\n",
        "from IPython.display import display, update_display\n",
        "from math import log2\n",
        "import torch\n",
        "from min_dalle import MinDalle\n",
        "\n",
        "model = MinDalle(\n",
-        "    dtype=torch.float16,\n",
+        "    dtype=getattr(torch, dtype),\n",
        "    is_mega=True, \n",
        "    is_reusable=True\n",
        ")"