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)就可以正常运行了。。。。