微信公众号搜"智元新知"关注
微信扫一扫可直接关注哦!

如何检查是否在CUDA上建模?

如何解决如何检查是否在CUDA上建模?

我想检查模型是否在CUDA上。该怎么做?

import torch
import torchvision
model = torchvision.models.resnet18()
model.to('cuda')

表明model.is_cuda()无效。

解决方法

此代码应执行以下操作:

import torch
import torchvision
model = torchvision.models.resnet18()
model.to('cuda')
next(model.parameters()).is_cuda

出局:

True

请注意,is_cuda()中没有nn.Module方法。 另外请注意,model.to('cuda')model.cuda()相同,并且都在适当位置。

另一方面,移动data.to('cuda')的位置不正确,您通常会调用:

data = data.to('cuda')

将数据移至CUDA。

版权声明:本文内容由互联网用户自发贡献,该文观点与技术仅代表作者本人。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容, 请发送邮件至 dio@foxmail.com 举报,一经查实,本站将立刻删除。