OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

Add cost inference for MulGradient operator

Browse Source 
Summary: Add cost inference for MulGradient operator; also whitelist MulGradient in COMPUTE_OP_TYPES in dense_perf_estimation

Test Plan: buck run //caffe2/caffe2/python/operator_test:elementwise_ops_test

Reviewed By: CrazySherman

Differential Revision: D27614003

fbshipit-source-id: 30901e5e2b6ce7e2183c2362d1bf9f895046cf55
This commit is contained in:
Chao Kong 2021-04-15 15:53:07 -07:00 committed by Facebook GitHub Bot
parent 3fbca31be3
commit ff1498e668
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
caffe2/operators/elementwise_ops_schema.cc
View File

@ -320,7 +320,8 @@ OPERATOR_SCHEMA(MulGradient)
.NumInputs(3)
.NumOutputs(2)
.TensorInferenceFunction(ElementwiseGradientOpShapeInference)
.AllowInplace({{0, 0}, {0, 1}});
.AllowInplace({{0, 0}, {0, 1}})
.CostInferenceFunction(PointwiseCostInference<2>);
OPERATOR_SCHEMA(Div)
.NumInputs(2)