frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-21 05:34:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

Let tensor_a.new_tensor() be on tensor_a.device by default (#144958)

Browse Source 
Fixes #144957
Closes #73838 cc @albanD @ezyang

Currently, `tensor_a.new_tensor()` will return a on-cpu tensor no matter where is `tensor_a`. This differs from the document and is a side-effect of https://github.com/pytorch/pytorch/pull/41984.

See #144957 how current logic breaks dynamo.

This PR restore the documented behavior and add tests for `new_tensor`.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/144958
Approved by: https://github.com/ezyang
This commit is contained in:
Yichen Yan
2025-01-24 22:12:29 +00:00
committed by PyTorch MergeBot
parent 2a70de7e92
commit d4171b724e
2 changed files with 49 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
test/test_tensor_creation_ops.py
View File

@ -2758,6 +2758,22 @@ class TestTensorCreation(TestCase):
sparse_size, dtype=torch.float64)
self.assertEqual(sparse_with_dtype.device, torch.device('cpu'))
@onlyCUDA
@onlyNativeDeviceTypes
def test_new_tensor_device(self, device):
torch_device = torch.device(device)
cpu_device = torch.device('cpu')
tensor = torch.tensor((1, 2, 3), device=device)
# need more than one device_type to test this
assert self.device_type == 'cuda'
for left, right in product([tensor, tensor.cpu()], [tensor, tensor.cpu()]):
for device_arg in [torch_device, cpu_device, None]:
if device_arg is None:
self.assertEqual(left.new_tensor(right).device, left.device)
else:
self.assertEqual(left.new_tensor(right, device=device_arg).device, device_arg)
def _test_signal_window_functions(self, name, dtype, device, **kwargs):
import scipy.signal as signal