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

Revert "Add torch.serialization.skip_data context manager (#134504)"

Browse Source 
This reverts commit 94db935749b8de99d8c3ab23fb880c67c8f3e67a.

Reverted https://github.com/pytorch/pytorch/pull/134504 on behalf of https://github.com/kit1980 due to See D62082697 ([comment](https://github.com/pytorch/pytorch/pull/134504#issuecomment-2327542276))
This commit is contained in:
PyTorch MergeBot
2024-09-03 22:21:27 +00:00
parent 85fa019697
commit 2fd36086bc
7 changed files with 26 additions and 261 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
torch/_utils.py
View File

@ -3,6 +3,7 @@ import copyreg
import functools
import logging
import sys
import threading
import traceback
import warnings
from collections import defaultdict
@ -108,13 +109,16 @@ def _get_async_or_non_blocking(function_name, non_blocking, kwargs):
return kwargs["async"]
_thread_local_state = threading.local()
def _get_restore_location(device):
"""Return the map_location location.
Used for rebuild functions where the tensor device is distinct from the storage
"""
map_location = torch.serialization._serialization_tls.map_location
map_location = getattr(_thread_local_state, "map_location", None)
if map_location is None:
return device
else: