pytorch
fca15236 - implement numel and tests for nested tensor (#80424)

Commit
2 years ago
implement numel and tests for nested tensor (#80424) Add numel implementation for Nested Tensor. Currently the construction of nested size and nested_strides assume contiguous. This implementation was based off of the safe_compute_numel(). Having a TORCH_CHECK in a for loop kinda feels bad but I don't really know how performant numel needs to be. Since nested size is stored as a tensor: `nested_size_tensor().cumprod(dim=1).sum(dim=0)[1].item() ` Would also get the job done. Pull Request resolved: https://github.com/pytorch/pytorch/pull/80424 Approved by: https://github.com/cpuhrsch
Author
Committer
Parents
Loading