pytorch
7cfbc858 - [fx_acc] [fx2trt] add acc op mapper for argmin and converter for topk (#63823)

Commit
3 years ago
[fx_acc] [fx2trt] add acc op mapper for argmin and converter for topk (#63823) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/63823 Add mapper for `torch.argmin` which maps it to `acc_ops.flatten` (optional) + `acc_ops.topk` + `acc_ops.getitem` + `acc_ops.squeeze` (optional). This diff doesn't allow mapping if `dim=None && keepdim=True` in `torch.argmin`. Add fx2trt converter for `acc_ops.topk`. Test Plan: buck test mode/opt glow/fb/fx/oss_acc_tracer:test_acc_tracer -- test_argmin buck run mode/opt caffe2/torch/fb/fx2trt:test_topk Reviewed By: jfix71 Differential Revision: D30501771 fbshipit-source-id: 0babc45e69bac5e61ff0b9b4dfb98940398e3e57
Author
Parents
Loading