Tensor moving from initial device to default device when setting in place

1

When pointing a new tensor to an existing tensor storage and trying to allocate more space than the original tensor, the storage is moved to gpu 0.

>>> import torch
>>> torch.__version__
'0.4.1'
>>> a = torch.rand(2, device=1)
>>> a
tensor([0.3587, 0.7057], device='cuda:1')
>>> b = a.new().set_(a.data.storage(), 0, (2,))
>>> b
tensor([0.3587, 0.7057], device='cuda:1')
>>> c = a.new().set_(a.data.storage(), 0, (3,))
>>> c
tensor([0.3587, 0.7057, 0.0000], device='cuda:0')
>>> a
tensor([0.3587, 0.7057], device='cuda:0')

b stayed on gpu 1 but when creating c which requested more memory than the size of a, the storage was moved to gpu 0 - which is the current context by default.

The workaround is to perform the operations within the expected context:

>>> a = torch.rand(2, device=1)
>>> with torch.cuda.device(a.device.index):
...     c = a.new().set_(a.data.storage(), 0, (3,))
>>> c
tensor([0.7182, 0.3534, 0.0000], device='cuda:1')

That behavior was introduced starting with version 0.4.1 and was not present in 0.4.0
Shouldn’t the expected behavior be for the storage to remain on the same device even if more memory is being requested and without having to re-specify the context?

1 Like