CUDA error: out of memory (allocate at /pytorch/aten/src/THC/THCCarchingAllocator.cpp:510)

博客主要讲述了在PyTorch中,多GPU训练后单张GPU加载模型会出现报错的问题及解决办法。当整个model打包保存时,加载模型需按特定步骤操作,如指定路径、加载检查点、处理DataParallel等,最后使模型进入评估模式。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

在pytorch中,使用多GPU训练,但测试加载模型时使用单张GPU会出现这类报错

解决:

整个model打包保存的情况下,加载模型时应:

model_path = 'model.tar'

checkpoint = torch.load(model_path)

model = checkpoint['model']

if isintance(model, torvh.nn.DataParallel):

      model = model.module

model = model.cuda()

model.eval()

 

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值