You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
19 lines
315 B
19 lines
315 B
1 month ago
|
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()
|