nn.Embedding踩坑:RuntimeError: Expected object of device type cuda but got device type cpu

        最近,我在做MAE相关代码实验的时候,网络的构架中有用到nn.Embedding()这个函数,刚开始,我用cpu跑训练的时候是没有报错的,但是cpu跑太慢了,我就改成gpu,结果出现了如下错误:

'''改->'''
device = torch.device("cpu")->device = torch.device("cuda")
'''结果'''
mask_tokens = mask_tokens + self.decoder_pos_emb(mask_index)
...
return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)

RuntimeError: Expected object of device type cuda but got 
device type cpu for argument #3 'index' in call to _th_index_select

  根据错误提示,我定位到了mask_tokens = mask_tokens+self.decoder_pos_emb(mask_index)这段代码,并且是self.decoder_pos_emb(mask_index)这个函数错误,因为只有self.decoder_pos_emb()这个函数涉及到Embedding():

self.decoder_pos_emb(mask_index)

self.decoder_pos_emb = nn.Embedding(self.encoder.dec_total_num, decoder_dim)
#构造一个可学习的参数,如果括号中的参数为(64,512),那么就构造一个64行,512列的tensor矩阵

mask_index = mask_index.type(torch.LongTensor)
#mask_index是需要对Embedding矩阵进行掩码的矩阵序列,并将其转化成LongTensor-int64位
#(Embedding函数要求)

        我觉得self.decoder_pos_emb(mask_index)应该是没错的,所以应该是传入的参数mask_index出错,所以我出查了查torch.LongTensor参数问题。查了半天,终于找到了问题所在,原来cpu->torch.LongTensor,  gpu->torch.cuda.LongTensor,具体可查找官网:

torch.Tensor — PyTorch 2.0 documentation

nn.Embedding踩坑:RuntimeError: Expected object of device type cuda but got device type cpu_第1张图片

         所以只要把数据类型torch.LongTensor改成torch.cuda.LongTensor就没有问题了:

mask_index = mask_index.type(torch.cuda.LongTensor)

你可能感兴趣的:(python,pytorch,深度学习)