Fix SmoothL1Loss op sigma attr problem.

This commit is contained in:
liangchenghui 2020-10-29 15:17:30 +08:00
parent d8a3ec4ead
commit 9d0d528e99
1 changed files with 1 additions and 0 deletions

View File

@ -1869,6 +1869,7 @@ class SmoothL1Loss(PrimitiveWithInfer):
validator.check_value_type('beta', beta, [float], self.name)
validator.check('beta', beta, '', 0, Rel.GT, self.name)
self.init_prim_io_names(inputs=['prediction', 'target'], outputs=['output'])
self.add_prim_attr('sigma', beta)
def infer_shape(self, prediction, target):
validator.check('prediction shape', prediction, 'target shape', target, Rel.EQ, self.name)