[Benchmark] Benchmark to(device) (#385)

pytorch · May 18, 2023 · a84add9 · a84add9
1 parent da83c3b
commit a84add9
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/benchmarks/common/common_ops_test.py b/benchmarks/common/common_ops_test.py
@@ -168,6 +168,22 @@ def exec_select():
     benchmark.pedantic(exec_select, iterations=100, rounds=100)
 
 
+@pytest.mark.skipif(not torch.cuda.device_count(), reason="No cuda device")
+def test_to(benchmark, td):
+    benchmark.pedantic(td.to, args=("cuda:0",), iterations=100, rounds=1000)
+
+
+@pytest.mark.skipif(not torch.cuda.device_count(), reason="No cuda device")
+def test_to_nonblocking(benchmark, td):
+    benchmark.pedantic(
+        td.to,
+        args=("cuda:0",),
+        kwargs={"non_blocking", True},
+        iterations=100,
+        rounds=1000,
+    )
+
+
 def main():
     # creation
     td = TensorDict({}, [3, 4])