Hi batch_size = 100 train = TensorDataset(train_X,train_Y) trainLoader = DataLoader(train,batch_size, shuffle=True,num_workers=1) model=model.cuda() for i,data in enumerate(trainLoader): train_x, train_y = data train_x = train_x/255 #train_x,train_y = Variable(train_x), Varia…

Are you using the same batch size to compute the new output? If you don’t need to train the model, wrap the code in a with torch.no_grad() block to save memory.

GPU gets out of memory

Eta_C January 9, 2020, 1:48am 3

You may be interested in this topic. There are some discussion here.

1 Like