From 743b36f0ce3103016c1aac396c29e2562fee3f95 Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Wed, 17 Jan 2024 16:44:15 +0000
Subject: [PATCH] hotfix: copy size is in bytes

---
 extra/transfer_speed.py | 18 ++++++++++++++++++
 tinygrad/device.py      |  2 +-
 2 files changed, 19 insertions(+), 1 deletion(-)
 create mode 100644 extra/transfer_speed.py

diff --git a/extra/transfer_speed.py b/extra/transfer_speed.py
new file mode 100644
index 00000000..9abe475a
--- /dev/null
+++ b/extra/transfer_speed.py
@@ -0,0 +1,18 @@
+from tinygrad import Tensor, Device
+
+#N = 1024
+N = 32
+t = Tensor.rand(N, N, N, device="CPU").realize()
+d1 = Device.DEFAULT + ":1"
+d2 = Device.DEFAULT + ":2"
+d3 = Device.DEFAULT + ":3"
+
+for i in range(3):
+  t.to_(d1)
+  t.realize()
+  # t.to_("CPU")
+  # t.realize()
+  t.to_(d2)
+  t.realize()
+  t.to_(d3)
+  t.realize()
diff --git a/tinygrad/device.py b/tinygrad/device.py
index aa5ea9ca..1f636764 100644
--- a/tinygrad/device.py
+++ b/tinygrad/device.py
@@ -130,7 +130,7 @@ class _BufferCopy(JITRunner):
     if wait or DEBUG >= 2:
       Device[dest.device].synchronize()
       et = time.perf_counter() - st
-    update_stats(colored(f"copy {dest.size:8d}, {dest.device[:7]:>7s} <- {src.device[:7]:7s}", "yellow"),
+    update_stats(colored(f"copy {dest.size*dest.dtype.itemsize:8d}, {dest.device[:7]:>7s} <- {src.device[:7]:7s}", "yellow"),
                  0, dest.size*dest.dtype.itemsize, {}, et, 2, jit, device=dest.device)
 BufferCopy = _BufferCopy()