Issue
the code
c = torch.rand((2000, 64, 64)).to('cuda')
d = torch.rand((2000, 64, 64)).to('cuda')
t3 = time.time()
s1 = c+d
s2 = torch.concat((a, b), dim=2)
t4 = time.time()
s1's device is gpu, but s2's device is cpu.
So I can't understand it. What is the principle of this?
Solution
Torch will do an operation if all the necessary variables for the operation are on the same device.
In your example, we can assume that a
and b
were on the CPU. Hence, torch.concat((a, b), dim=2)
is on the CPU as well. When you write .to('cuda')
, you move c
and d
to the GPU, thus s1
is on the GPU too.
Answered By - PlainRavioli
0 comments:
Post a Comment
Note: Only a member of this blog may post a comment.