Bug fixes

m-gopichand · Aug 20, 2024 · 01f04ff · 01f04ff
1 parent d753d06
commit 01f04ff
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/nbs/03_convolutions.ipynb b/nbs/03_convolutions.ipynb
@@ -1525,6 +1525,9 @@
     "#|export\n",
     "def_device = 'mps' if torch.backends.mps.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu'\n",
     "\n",
+    "def device():\n",
+    "    return 'mps' if torch.backends.mps.is_available() else 'cuda' if torch.cuda.is_available() else 'cpu'\n",
+    " \n",
     "def to_device(x, device=def_device):\n",
     "    if isinstance(x, torch.Tensor): return x.to(device)\n",
     "    if isinstance(x, Mapping): return {k:v.to(device) for k,v in x.items()}\n",