pytorch1.6 使用在含有RNN/LSTM网络中使用amp,出错:CUDNN_STATUS_BAD_PARAM,

 with torch.cuda.amp.autocast():
        preds = model(inp)
        batch_size = inp.size(0)
        text, length = converter.encode(labels)
        preds_size = torch.IntTensor([preds.size(0)] * batch_size)
        loss = criterion(preds, text, preds_size, length)
    optimizer.zero_grad()
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

上面是按照官方的代码使用 amp,报错位置在我的rnn部分

然后在forward里面打印了一下,rnn之前输出,也就是rnn即将输入部分。发现使用amp网络里面自动变成了float16,但是rnn是需要float32的,所以我在rnn之前做了一个input=input.to(torch.float32)就可以正常运行了。。。。

你可能感兴趣的:(pytorch)