Skip to content

Commit 965d658

Browse files
authored
Module support H2D move (#1831)
1 parent 6285685 commit 965d658

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

mindnlp/core/nn/modules/module.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -688,7 +688,10 @@ def __dir__(self):
688688
return sorted(keys)
689689

690690
def cuda(self):
691-
return self
691+
return self._apply(lambda t: t.move_to('GPU'))
692+
693+
def npu(self):
694+
return self._apply(lambda t: t.move_to('Ascend'))
692695

693696
def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
694697
missing_keys, unexpected_keys, error_msgs):

0 commit comments

Comments
 (0)