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

Restore fake device (#157972)

Browse Source 
Fixes #ISSUE_NUMBER

Pull Request resolved: https://github.com/pytorch/pytorch/pull/157972
Approved by: https://github.com/ezyang
This commit is contained in:
angelayi
2025-07-11 16:11:57 +00:00
committed by PyTorch MergeBot
parent 27c50799c1
commit 1f1f22991d
2 changed files with 15 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
test/test_fake_tensor.py
View File

@ -253,6 +253,19 @@ class FakeTensorTest(TestCase):
assert x.copy_(y).device.type == "cpu"
assert y.copy_(x).device.type == "cuda"
def test_fake_device(self):
t = torch.ones(3)
t = t.view(1, 3)
fake_mode1 = FakeTensorMode(allow_non_fake_inputs=True)
fake_t = fake_mode1.from_tensor(t)
fake_t.fake_device = torch.device("cuda")
fake_mode2 = FakeTensorMode(allow_non_fake_inputs=True)
new_fake_t = fake_mode2.from_tensor(fake_t)
self.assertEqual(new_fake_t.device, fake_t.device)
def test_fake_dispatch_keys(self):
with FakeTensorMode():
x = torch.rand([4])