added gpu to extreme classification

hqman · Jan 22, 2023 · 37f62b2 · 37f62b2
1 parent a0bc230
commit 37f62b2
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/analytics-and-ml/extreme-classification/extreme-classification.ipynb b/analytics-and-ml/extreme-classification/extreme-classification.ipynb
@@ -388,9 +388,15 @@
    ],
    "source": [
     "from sentence_transformers import SentenceTransformer\n",
+    "import torch\n",
+    "\n",
+    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
     "\n",
     "# load the model from huggingface\n",
-    "model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')\n",
+    "model = SentenceTransformer(\n",
+    "    'sentence-transformers/all-MiniLM-L6-v2',\n",
+    "    device=device\n",
+    ")\n",
     "model"
    ]
   },