10000 Allow BUILD/NEWOBJ instruction for items added via torch.serialization.add_safe_globals by mikaylagawarecki · Pull Request #129251 · pytorch/pytorch · GitHub
[go: up one dir, main page]
Skip to content

Allow BUILD/NEWOBJ instruction for items added via torch.serialization.add_safe_globals #129251

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 7 commits into from
Closed

Allow BUILD/NEWOBJ instruction for items added via torch.serialization.add_safe_globals #129251

Changes from 1 commit
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
9611156
Update
mikaylagawarecki Jun 21, 2024
4d5f7c3
Update
mikaylagawarecki Jun 21, 2024
bf878be
Update
mikaylagawarecki Jun 21, 2024
b4c10b7
Update
mikaylagawarecki Jun 21, 2024
06df8ce
Update
mikaylagawarecki Jun 21, 2024
cd4dbc9
Update
mikaylagawarecki Jun 24, 2024
5c9cdb3
Update
mikaylagawarecki Jun 24, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Update 
[ghstack-poisoned]
  • Loading branch information
@mikaylagawarecki
mikaylagawarecki committed Jun 21, 2024
commit 06df8ce71b0d17d0102adac6cc180986da3a6454
10 changes: 10 additions & 0 deletions test/distributed/_tensor/test_dtensor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -536,6 +536,16 @@ def test_dtensor_save_load(self):
buffer.seek(0)
reloaded_st = torch.load(buffer)
self.assertEqual(sharded_tensor, reloaded_st)
# Test weights_only load
try:
torch.serialization.add_safe_globals(
[DTensor, DeviceMesh, Shard, DTensorSpec, TensorMeta]
)
buffer.seek(0)
reloaded_st = torch.load(buffer, weights_only=True)
self.assertEqual(sharded_tensor, reloaded_st)
finally:
torch.serialization.clear_safe_globals()


class DTensorMeshTest(DTensorTestBase):
Expand Down
Loading
0