Module support H2D move (#1831)

mindspore-lab · Nov 21, 2024 · 965d658 · 965d658
1 parent 6285685
commit 965d658
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/mindnlp/core/nn/modules/module.py b/mindnlp/core/nn/modules/module.py
@@ -688,7 +688,10 @@ def __dir__(self):
         return sorted(keys)
 
     def cuda(self):
-        return self
+        return self._apply(lambda t: t.move_to('GPU'))
+
+    def npu(self):
+        return self._apply(lambda t: t.move_to('Ascend'))
 
     def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
                               missing_keys, unexpected_keys, error_msgs):