Fix FLOPs measurement error on CUDA

Author: LukasHedegaard

CHANGELOG.md

@@ -12,6 +12,7 @@ From v1.0.0 and on, the project will adherence strictly to Semantic Versioning.
 ## [0.3.5] - 2023-08-08
 ### Fixed
 - Updated on-gpu model benchmaking with best-practices on `cuda.Event` and `cuda.synchronize`.
+- FLOPs measurement error on CUDA.
 
 
 ## [0.3.4] - 2022-02-22

pytorch_benchmark/benchmark.py

@@ -343,7 +343,11 @@ def benchmark(
         batch_size=1,
     )
 
-    flops = measure_flops(model, sample1, print_details)
+    with torch.no_grad():
+        flops = measure_flops(
+            model, transfer_to_device_fn(sample1, model_device), print_details
+        )
+
     if _is_valid(flops):
         results["flops"] = flops
         print_fn(f"Model FLOPs: {flops} ({format_num(flops)})")