Update

[ghstack-poisoned]
pytorch · Oct 18, 2024 · 26affa0 · 26affa0
1 parent 9c7ed32
commit 26affa0
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/benchmarks/common/h2d_test.py b/benchmarks/common/h2d_test.py
@@ -99,12 +99,12 @@ def test_to(
         tensordict_logger.info(f"td size {td.bytes() / 1024 / 1024:.2f} Mb")
         pin_mem = default_device.type == "cuda"
         if consolidated is True:
-            td = td.consolidate(pin_memory=pin_mem, set_on_tensor=True)
+            td = td.consolidate(pin_memory=pin_mem)
 
         if consolidated == "within":
 
             def to(td, num_threads):
-                return td.consolidate(pin_memory=pin_mem, set_on_tensor=True).to(
+                return td.consolidate(pin_memory=pin_mem).to(
                     default_device, num_threads=num_threads
                 )
 
@@ -127,12 +127,12 @@ def test_to_njt(
         tensordict_logger.info(f"njtd size {njt_td.bytes() / 1024 / 1024 :.2f} Mb")
         pin_mem = default_device.type == "cuda"
         if consolidated is True:
-            njt_td = njt_td.consolidate(pin_memory=pin_mem, set_on_tensor=True)
+            njt_td = njt_td.consolidate(pin_memory=pin_mem)
 
         if consolidated == "within":
 
             def to(td, num_threads):
-                return td.consolidate(pin_memory=pin_mem, set_on_tensor=True).to(
+                return td.consolidate(pin_memory=pin_mem).to(
                     default_device, num_threads=num_threads
                 )