issue/890 - 为nn.module添加to函数 #891
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
pengcheng888
force-pushed
the
issue/890
branch
from
78c0706 to
dfb1e8d
Compare
pengcheng888
commented
Jan 7, 2026
| if (param.shape == input_param.shape) and ( | ||
| param.dtype == input_param.dtype | ||
| ): | ||
| param.copy_(input_param) |
Collaborator
Author
There was a problem hiding this comment.
因为，infinicore的两个tensor之间的copy操作，是支持从 cpu直接拷贝到gpu的。
现在的权重加载判断是： 模型weight和权重文件，二者shape和dtype同一样时，可以拷贝数据，否则报错。
pengcheng888
commented
Jan 7, 2026
| pass | ||
| raise KeyError("not support") | ||
|
|
||
| def _apply(self, fn, recurse=True): |
Collaborator
Author
There was a problem hiding this comment.
to函数的部分参考了torch的写法
pengcheng888
commented
Jan 7, 2026
| raise KeyError("not support") | ||
| for key, param in self._parameters.items(): | ||
| if param is not None: | ||
| setattr(self, key, fn(param)) |
Collaborator
Author
There was a problem hiding this comment.
使用=符号，赋值不成功，不知为何。
Collaborator
Author
There was a problem hiding this comment.
最后用了setattr(self, key, fn(param))
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
目标版本
main
功能描述
python的nn.Module中添加 to函数：
实现将model权重to到gpu上的操作
model.to("cuda")
测试结果
该功能的测试脚本，测试通过
修改后infinilm中的llama模型，测试通过
