梯度更新流程

loss = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

train_pred = model(data[0].cuda())
train_loss = loss(train_label, data[1].cuda())
train_loss.backward()
optimizer.step()

 

来源：https://www.cnblogs.com/JunzhaoLiang/p/11869475.html