the result from torch.concat() is stored in cpu(memory)?
Question:
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?
Answers:
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.
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?
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.