fix gamma on CPU and GPU

mindspore-lab · Oct 30, 2024 · ecdd3db · ecdd3db
1 parent e9db14d
commit ecdd3db
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)