[Pytorch 꿀팁] Cuda:out of memory

딥러닝 코드를 작성하다 보면 Cuda: out of memory를 볼 때가 있다.

내 코드가 그렇게 거대하지 않은데도 불구하고 문제가 발생한다면

아래 부분을 확인해 볼 필요가 있다.

for epoch in range(EPOCHS):
    loss_val_sum = 0
    for batch_in,batch_out in train_iter:
        # Forward path
        y_pred = C.forward(batch_in.view(-1,1,28,28).to(device))
        loss_out = loss(y_pred,batch_out.to(device))
        # Update
        optm.zero_grad()   
        loss_out.backward()   
        optm.step()    
        ################################################
        loss_val_sum += loss_out
        ################################################
    loss_val_avg = loss_val_sum/len(train_iter)
    # Print
    if ((epoch%print_every)==0) or (epoch==(EPOCHS-1)):
        train_accr = func_eval(C,train_iter,device)
        test_accr = func_eval(C,test_iter,device)
        print ("epoch:[%d] loss:[%.3f] train_accr:[%.3f] test_accr:[%.3f]."%
               (epoch,loss_val_avg,train_accr,test_accr))

잘 보면 loss_val_sum에 loss_out.item()이 아닌 loss_out을 그대로 더하고 있다.
loss.item()은 loss의 value of tensor값을, loss는 loss 객체 자체를 반환하기 때문에 메모리를 과하게 차지하고 있을 수 있다.