使用PyTorch-OpCounter工具出现错误：module must have its parameters and buffers on device cuda

最近使用PyTorch-OpCounter工具来计算网络的FLOPs出现以下错误，在网上搜寻了很久都没有答案，后面通过分析解决问题，记录一下。
出现的错误：module must have its parameters and buffers on device cuda:0(device_ids[0]) but found one of them on device:cpu

1、背景
首先是正常训练好一个模型，然后写了一个脚本使用OpCounter工具来计算FLOPs。当使用脚本的时候一直出现这个错误，没法计算出FLOPs。后来查询资料，DataParrallel的forward源码中规定parameters和buffers必须在一块显卡上，一般是GPU0。
2、分析
首先使用next(model.buffers()).is_cuda和next(model.parameters ()).is_cuda来检查buffers和parameters 是否真的不在同个GPU上，结果输出均是True，两者均在GPU上，并不是错误提示所说的其中一个在CPU上。
怀疑是在利用OpCounter工具来计算的时候，其中一个并不在GPU上。后来想到估计问题是出在使用了一机多卡训练模型导致其中某个过程是在CPU上，所以将模型变成单卡来测试。

#OpCounter工具计算FLOPs方法
from torchvision.models import resnet18
from thop import profile
import torch
model = resnet18()
input = torch.randn(1, 3, 224, 224)
macs, params = profile(model, inputs=(input, )) 
print(macs)

3、解决
将模型提取到单GPU，不再使用nn.DataParallel()。最后测试成功。

new_state_dic=OrderedDict()
for k,v in model_state_dict.items():
    name=k[7:]
    new_state_dic[name]=v
model.load_state_dict(new_state_dict)
model=model.cuda()