amend

Author: Vincent Moens

.ci/docker/requirements.txt

@@ -30,7 +30,7 @@ pytorch-lightning
 torchx
 # TODO: use stable 0.5 when released
 -e git+https://github.com/pytorch/rl.git#egg=torchrl
--e git+https://github.com/pytorch/tensordict.git#egg=tensordict
+-e git+https://github.com/pytorch/tensordict.git@del-futures#egg=tensordict
 ax-platform
 nbformat>==5.9.2
 datasets

intermediate_source/pinmem_nonblock.py

@@ -500,18 +500,13 @@ def pin_copy_to_device_nonblocking(*tensors):
 from torch.utils.benchmark import Timer
 
 # Create the dataset
-for s0 in (10, 100, 1000, 10_000, 100_000, 1_000_000):
-    for s1 in (10, 100, 1000, 10_000, 100_000, 1_000_000):
-        if s0 * s1 >= 1e9:
-            continue
-        print(f"\n\ns0={s0}, s1={s1}")
-        td = TensorDict({str(i): torch.randn(s0) for i in range(s1)})
-
-        # Runtimes
-        copy_blocking = timer("td.to('cuda:0', non_blocking=False)")
-        copy_non_blocking = timer("td.to('cuda:0')")
-        copy_pin_nb = timer("td.to('cuda:0', non_blocking_pin=True, num_threads=0)")
-        copy_pin_multithread_nb = timer("td.to('cuda:0', non_blocking_pin=True, num_threads=4)")
+td = TensorDict({str(i): torch.randn(1_000_000) for i in range(1000)})
+
+# Runtimes
+copy_blocking = timer("td.to('cuda:0', non_blocking=False)")
+copy_non_blocking = timer("td.to('cuda:0')")
+copy_pin_nb = timer("td.to('cuda:0', non_blocking_pin=True, num_threads=0)")
+copy_pin_multithread_nb = timer("td.to('cuda:0', non_blocking_pin=True, num_threads=4)")
 
 # Rations
 r1 = copy_non_blocking / copy_blocking