Quick fix scalar add flaws in MindSpore

5 years ago · 1950002e83
parent b2cff2842d
commit 1950002e83
1 changed files with 2 additions and 1 deletions
--- a/mindspore/nn/probability/bijector/scalar_affine.py
+++ b/mindspore/nn/probability/bijector/scalar_affine.py
@ -69,6 +69,7 @@ class ScalarAffine(Bijector):
            param=param)

        self.abs = P.Abs()
+        self.oneslike = P.OnesLike()
        self.log = log_generic

    @property
@ -92,7 +93,7 @@ class ScalarAffine(Bijector):
            f(x) = a * x + b
        """
        x = self._check_value(x, 'value')
-        return self.scale * x + self.shift
+        return self.scale * x + self.shift * self.oneslike(x)

    def _inverse(self, y):
        r"""