CODEWITHSUNDEEP
nlppytorchhuggingface-transformers
Ali Ebrahimi
Ali Ebrahimi

Reputation: 29

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?

Upvotes: 1

Views: 285

Answers (1)

cronoik
cronoik

Reputation: 19495

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()
})

Upvotes: 2

Related Questions