From 1950002e83aec24901f295bdb60b38e2da06d8ad Mon Sep 17 00:00:00 2001
From: peixu_ren <peixu.ren1@huawei.com>
Date: Tue, 25 Aug 2020 23:17:59 -0400
Subject: [PATCH] Quick fix scalar add flaws in MindSpore

---
 mindspore/nn/probability/bijector/scalar_affine.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mindspore/nn/probability/bijector/scalar_affine.py b/mindspore/nn/probability/bijector/scalar_affine.py
index b75298be2..e76518742 100644
--- 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"""
-- 
GitLab