Update

[ghstack-poisoned]
pytorch · Oct 17, 2024 · a9dbfb3 · a9dbfb3
1 parent 26642ca
commit a9dbfb3
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/benchmarks/common/h2d_test.py b/benchmarks/common/h2d_test.py
@@ -76,7 +76,9 @@ class TestTo:
     def test_to(
         self, benchmark, consolidated, td, default_device, compile_mode, num_threads
     ):
-        tensordict_logger.info(f"td size (Gb) {td.bytes() / 1024 / 1024 / 1024 :.2f} GB")
+        tensordict_logger.info(
+            f"td size {td.bytes() / 1024 / 1024:.2f} Mb"
+        )
         if consolidated:
             td = td.consolidate()
 
@@ -94,7 +96,9 @@ def to(td, num_threads):
     def test_to_njt(
         self, benchmark, consolidated, njt_td, default_device, compile_mode, num_threads
     ):
-        tensordict_logger.info(f"njtd size (Gb) {td.bytes() / 1024 / 1024 / 1024 :.2f} GB")
+        tensordict_logger.info(
+            f"njtd size {td.bytes() / 1024 / 1024 :.2f} Mb"
+        )
         if consolidated:
             njt_td = njt_td.consolidate()