topk tensor k support (#39407)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/39407
- support passing a single element tensor as k for topk module
- support passing a single element tensor to constant fill output
Test Plan:
buck test dper3/dper3/modules/tests:core_modules_test -- test_topk_gating_without_split_examples_tensor_k
buck test caffe2/caffe2/python:hypothesis_test -- test_constant_fill_from_tensor
Reviewed By: huayuli00
Differential Revision: D21843739
fbshipit-source-id: 0c5f5c03e9f57eeba40c0068784625164c2527ec