update compute

Author: nathanhubens

fasterbench/compute.py

@@ -41,6 +41,14 @@ def as_dict(self) -> Dict[str, float | str]:
 def compute_compute(model: nn.Module, dummy_input: torch.Tensor) -> ComputeMetrics:
     """Return **ComputeMetrics** for a single forward pass."""
 
+    try:
+        model_device = next(model.parameters()).device
+    except StopIteration:
+        model_device = torch.device("cpu")  # model has no parameters
+
+    if dummy_input.device != model_device:
+        dummy_input = dummy_input.to(model_device)
+
     params_m = round(
         sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6, 3
     )

nbs/03_compute.ipynb

@@ -67,6 +67,14 @@
     "def compute_compute(model: nn.Module, dummy_input: torch.Tensor) -> ComputeMetrics:\n",
     "    \"\"\"Return **ComputeMetrics** for a single forward pass.\"\"\"\n",
     "\n",
+    "    try:\n",
+    "        model_device = next(model.parameters()).device\n",
+    "    except StopIteration:\n",
+    "        model_device = torch.device(\"cpu\")  # model has no parameters\n",
+    "\n",
+    "    if dummy_input.device != model_device:\n",
+    "        dummy_input = dummy_input.to(model_device)\n",
+    "\n",
     "    params_m = round(\n",
     "        sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6, 3\n",
     "    )\n",
@@ -99,7 +107,13 @@
    "source": []
   }
  ],
- "metadata": {},
+ "metadata": {
+  "kernelspec": {
+   "display_name": "python3",
+   "language": "python",
+   "name": "python3"
+  }
+ },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }