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

Support is_mtia attribute. (#108307) (#108310)

Browse Source 
Summary:

FBGEMM uses `self.iter.is_cuda` to check if the tensor is for CUDA. This diff enables similar feature `self.iter.is_mtia` for tensors with MTIA device key.

Test Plan: See diff D48693225

Reviewed By: jackm321

Differential Revision: D48809191

Pull Request resolved: https://github.com/pytorch/pytorch/pull/108310
Approved by: https://github.com/albanD
This commit is contained in:
Jun Luo
2023-09-01 01:25:36 +00:00
committed by PyTorch MergeBot
parent d569e506ab
commit 8289ad8e5e
6 changed files with 23 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
torch/csrc/jit/frontend/sugared_value.cpp
View File

@ -120,6 +120,7 @@ std::shared_ptr<SugaredValue> SimpleValue::attr(
{"is_sparse_csr", "prim"},
{"is_mkldnn", "prim"},
{"is_mps", "prim"},
{"is_mtia", "prim"},
{"is_quantized", "prim"},
{"is_vulkan", "prim"},
{"is_ipu", "prim"},