fix gamma on CPU and GPU (#1791)

mindspore-lab · Oct 30, 2024 · 47f2de8 · 47f2de8
1 parent e9db14d
commit 47f2de8
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/mindnlp/core/ops/random.py b/mindnlp/core/ops/random.py
@@ -1,4 +1,5 @@
 """random op"""
+import numpy as np
 import mindspore
 from mindspore import ops
 from mindspore.ops._primitive_cache import _get_cache_prim
@@ -101,4 +102,6 @@ def randperm(n, seed=0, offset=0, dtype=mindspore.int64):
     return randperm_op(mindspore.tensor([n]))
 
 def gamma(shape, alpha, beta):
+    if DEVICE_TARGET != 'Ascend':
+        return mindspore.tensor(np.random.gamma(alpha, 1/beta, shape))
     return ops.gamma(shape, alpha, beta)