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

[MTIA] Support loading Tensors on mtia:0 for pytorch code (#149327)

Browse Source 
Summary: The diff includes updates to the PyTorch code to enable loading tensors to MTIA.

Reviewed By: PatriceVignola

Differential Revision: D71176848

Pull Request resolved: https://github.com/pytorch/pytorch/pull/149327
Approved by: https://github.com/ezyang
This commit is contained in:
Sergey Zimin
2025-03-20 18:05:15 +00:00
committed by PyTorch MergeBot
parent 7bb9c36784
commit b99fc9d29f
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
torch/csrc/jit/serialization/unpickler.cpp
View File

@ -623,7 +623,8 @@ PickleOpCode Unpickler::readInstruction() {
}
if (device.is_cuda() || device.is_xpu() || device.is_meta() ||
device.is_hpu() || device.is_mps() || device.is_privateuseone()) {
device.is_mtia() || device.is_hpu() || device.is_mps() ||
device.is_privateuseone()) {
tensor = tensor.to(device, tensor.scalar_type());
} else if (device.type() != DeviceType::CPU) {
TORCH_CHECK(