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

Update pin memory related APIs to not pass 'device' argument (#131858)

Browse Source 
Based on https://github.com/pytorch/pytorch/pull/126376, this PR tries to update all PT callers (e.g., `Tensor.is_pinned()`, `Tensor.pin_memory()`) to not pass `device` argument.
As for `storage/untyped_storage.is_pinned()/pin_memory()`, we keep the `device` argument but passing `device` is discouraged. And if not given, the default `device` is still 'cuda' for BC.
Additionally, based on device-agnostic pin_memory, `pin_memory_device` argument of `torch.utils.data.DataLoader` is discouraged  now. For BC, explictly passing this argument is still effective. If not given, the default `device` will be the current accelerator.

Fixes #124908
Relates https://github.com/pytorch/pytorch/pull/126376

Pull Request resolved: https://github.com/pytorch/pytorch/pull/131858
Approved by: https://github.com/albanD

Co-authored-by: albanD <desmaison.alban@gmail.com>
This commit is contained in:
wizzniu
2025-01-15 17:23:33 +00:00
committed by PyTorch MergeBot
parent 0dca756832
commit c07dc64017
9 changed files with 61 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
torch/storage.py
View File

@ -359,7 +359,8 @@ class _StorageBase:
r"""Determine whether the CPU storage is already pinned on device.
Args:
device (str or torch.device): The device to pin memory on. Default: ``'cuda'``.
device (str or torch.device): The device to pin memory on (default: ``'cuda'``).
This argument is discouraged and subject to deprecated.
Returns:
A boolean variable.
@ -374,7 +375,8 @@ class _StorageBase:
r"""Copy the CPU storage to pinned memory, if it's not already pinned.
Args:
device (str or torch.device): The device to pin memory on. Default: ``'cuda'``.
device (str or torch.device): The device to pin memory on (default: ``'cuda'``).
This argument is discouraged and subject to deprecated.
Returns:
A pinned CPU storage.
@ -1158,7 +1160,8 @@ class TypedStorage:
r"""Determine whether the CPU TypedStorage is already pinned on device.
Args:
device (str or torch.device): The device to pin memory on. Default: ``'cuda'``
device (str or torch.device): The device to pin memory on (default: ``'cuda'``).
This argument is discouraged and subject to deprecated.
Returns:
A boolean variable.
@ -1170,7 +1173,8 @@ class TypedStorage:
r"""Copy the CPU TypedStorage to pinned memory, if it's not already pinned.
Args:
device (str or torch.device): The device to pin memory on. Default: ``'cuda'``.
device (str or torch.device): The device to pin memory on (default: ``'cuda'``).
This argument is discouraged and subject to deprecated.
Returns:
A pinned CPU storage.