refactor prim op for different backend

mindnlp/__init__.py

@@ -18,7 +18,6 @@
 """
 import os
 import platform
-from packaging import version
 
 # huggingface env
 if os.environ.get('HF_ENDPOINT', None) is None:
@@ -29,13 +28,15 @@
     del os.environ['RANK_TABLE_FILE']
 
 import mindspore
-from mindspore import context
 from mindspore._c_expression import MSContext # pylint: disable=no-name-in-module, import-error
 try:
     from mindspore._c_expression import disable_multi_thread
 except:
     disable_multi_thread = None
 
+if os.environ.get('DEVICE_TARGET', None) is not None:
+    mindspore.set_device(os.environ.get('DEVICE_TARGET'))
+
 # for different ascend devices
 if platform.system().lower() == 'linux':
     SOC = MSContext.get_instance().get_ascend_soc_version()

mindnlp/core/_C/__init__.py

@@ -198,7 +198,10 @@ def _step(self, step):
         Returns:
             Current seed and offset.
         """
-        return self._generator(STEP, (self._seed, self._offset, step,))[:2]
+        outs = self._generator(STEP, (self._seed, self._offset, step,))[:2]
+        for o in outs:
+            o._device = self.device
+        return outs
 
 default_generator = Generator()