Corrects oversights in device handling for GPU support

aehiggi · May 23, 2022 · d570932 · d570932
1 parent d3d9665
commit d570932
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/hw_openqa.ipynb b/hw_openqa.ipynb
@@ -613,7 +613,7 @@
     "         \n",
     "    \"\"\"\n",
     "    prompt_ids = eleuther_tokenizer(\n",
-    "        prompts, return_tensors=\"pt\", padding=True).input_ids\n",
+    "        prompts, return_tensors=\"pt\", padding=True).input_ids.to(device)\n",
     "        \n",
     "    with torch.inference_mode():\n",
     "        # Automatic mixed precision if possible.\n",
@@ -645,7 +645,7 @@
     "    for prompt, gen_id, gen_text, gen_prob in iterator:       \n",
     "        gen_tokens = eleuther_tokenizer.convert_ids_to_tokens(gen_id)\n",
     "        generated_text = gen_text[len(prompt): ]\n",
-    "        gen_prob = [float(x) for x in gen_prob.numpy()] # float for JSON storage\n",
+    "        gen_prob = [float(x) for x in gen_prob.cpu().numpy()] # float for JSON storage\n",
     "        ans_indices = _find_generated_answer(gen_tokens, newline=\"Ċ\")\n",
     "        answer_tokens = [gen_tokens[i] for i in ans_indices]\n",
     "        answer_probs = [gen_prob[i] for i in ans_indices]\n",