I tried the following code. When the code in try failed because of out of CUDA memory, I reduced the batch size to a half in except but it still appear the same issue for running the model in except but I'm sure half of the batch size is runnable since I have tried to directly run the code in except without trying the full batch. It works fine. By the way, is there any way to automatically set the batch size to fully use the CUDA memory without overflow?
try:
output = model(Variable(torch.LongTensor(np.array(x))).to(device),Variable(torch.LongTensor(np.array(pos))).to(device),Variable(torch.LongTensor(np.array(m))).to(device))
loss = criterion(output, Variable(torch.LongTensor(y)).to(device))#lb.transform(y)
loss.backward()
optimizer.step()
losses.append(loss.data.mean())
except:
half = int(len(x) / 2)
x1 = x[:half]
x2 = x[half:]
pos1 = pos[:half]
pos2 = pos[half:]
m1 = m[:half]
m2 = m[half:]
y1 = y[:half]
y2 = y[half:]
optimizer.zero_grad()
output = model(Variable(torch.LongTensor(np.array(x1))).to(device),Variable(torch.LongTensor(np.array(pos1))).to(device),Variable(torch.LongTensor(np.array(m1))).to(device))
loss = criterion(output, Variable(torch.LongTensor(y1)).to(device))#lb.transform(y)
loss.backward()
optimizer.step()
losses.append(loss.data.mean())
output = model(Variable(torch.LongTensor(np.array(x2))).to(device),Variable(torch.LongTensor(np.array(pos2))).to(device),Variable(torch.LongTensor(np.array(m2))).to(device))
loss = criterion(output, Variable(torch.LongTensor(y2)).to(device))#lb.transform(y)
loss.backward()
optimizer.step()
losses.append(loss.data.mean())