Thank you so much for your patient explaination. I also did as you told me, but still have memory problem. I strongly doubt that pytorch load all the dataset in gpu which are 12000 896 * 448 images (48G). But ImageNet has 124G and can run on 18 G gpus. I have 2 nodes which is 48 G for each node. I can’t even run it with batch size 1.
And how can I check if there are other tensors on gpu? Thank you.
else: # logging for classification
            #_, predicted = torch.max(out.data, 1) for classification
            total += targets.size(0)
            #correct += predicted.eq(targets.data).cpu().sum() for classification
            loss_total+=loss

            log_dict = {"iter": cur_iter, "loss": loss.item(), "epoch": epoch}
            train_log.write("{}\n".format(json.dumps(log_dict)))
            train_log.flush()