Method 1

torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True

Principle:

Cundnn follows the following criteria:

if the dimension or type of network input data changes little, set    torch.backends.cudnn.benchmark  =  true   It can increase the operation efficiency; If the input data of the network changes every iteration, cndnn will find the optimal configuration every time, which will improve the operation efficiency

 
Method 2

Tensor calculation should be written as follows:

train_loss += loss.item()