save model output in pytorch
dic = [] 
for step, batch in tqdm(enumerate(train_dataloader)):
inpt = batch[0].to(device)
msks = batch[1].to(device)
#Run the sentences through the model
outputs = model_obj(inpt, msks)
dic.append( {
    'hidden_states': outputs[2],
    'pooled_output': outputs[1]})

I want to save the model output in each iteration but I got the below error for a small set of datasets.
RuntimeError: CUDA out of memory.
notice that without the below code my model works correctly.
dic.append( { 'hidden_states': outputs[2], 'pooled_output': outputs[1]})
How can I save these outputs in each iteration?

First of all, you should always post the full error stacktrace. Secondly, you should move the outputs from your GPU when you want to store them to free up memory:
dic.append( {
    'hidden_states': outputs[2].detach().cpu().tolist(),
    'pooled_output': outputs[1].detach().cpu().tolist()
})

