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

Add meta device support to sparse compressed tensors (#120498)

Browse Source 
As in the title.

Unblocks https://github.com/pytorch/pytorch/pull/117907#discussion_r1499251745

Pull Request resolved: https://github.com/pytorch/pytorch/pull/120498
Approved by: https://github.com/ezyang
This commit is contained in:
Pearu Peterson
2024-02-24 22:40:45 +02:00
committed by PyTorch MergeBot
parent 834c7a1d3e
commit 5d71ba6885
10 changed files with 130 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
torchgen/model.py
View File

@ -84,6 +84,7 @@ class DispatchKey(Enum):
Sparse = auto()
SparseCsrCPU = auto()
SparseCsrCUDA = auto()
SparseCsrMeta = auto()
NestedTensor = auto()
Dense = auto()
@ -260,6 +261,7 @@ dispatch_keys = [
# kernels
DispatchKey.Meta,
DispatchKey.SparseMeta,
DispatchKey.SparseCsrMeta,
DispatchKey.QuantizedMeta,
DispatchKey.NestedTensorMeta,
DispatchKey.ZeroTensor,