Update as_safetensors+fp16.ipynb

as_safetensors+fp16.ipynb

@@ -107,6 +107,7 @@
         "import glob\n",
         "import torch\n",
         "import safetensors.torch\n",
+        "from functools import partial\n",
         "\n",
         "models = \"wd-1-4-anime_e1.ckpt, wd-1-4-anime_e1.yaml\" #@param {type:\"string\"}\n",
         "as_fp16 = True #@param {type:\"boolean\"}\n",
@@ -139,7 +140,7 @@
         "  elif (model_ext != \".safetensors\") & (model_ext != \".ckpt\"):\n",
         "    print(\"対応形式は.ckpt及び.safetensors並びに.yamlのみです\\n\" + f\"\\\"{model}\\\"は対応形式ではありません\")\n",
         "  else:\n",
-        "    load_model = safetensors.torch.load_file if model_ext == \".safetensors\" else torch.load\n",
+        "    load_model = partial(safetensors.torch.load_file, device=\"cpu\") if model_ext == \".safetensors\" else partial(torch.load, map_location=torch.device(\"cpu\"))\n",
         "    save_model = safetensors.torch.save_file if save_type == \".safetensors\" else torch.save\n",
         "    # convert model\n",
         "    with torch.no_grad():\n",