解决加载torch模型时出现CUDA out of memory

from transformers import BertForSequenceClassification, AdamW

pretrained = 'bert-base-chinese'

model = BertForSequenceClassification.from_pretrained(
    pretrained, 
    num_labels=121, 
    output_attentions = False, # 模型是否返回 attentions weights.
    output_hidden_states = False, # 模型是否返回所有隐层状态.
    return_dict=False
)
model.load_state_dict(torch.load("model/BERT_model.h5"))

在没有加载gpu的情况下为什么还会爆出CUDA内存不够的bug呢？
————————————————————————————
原因
保存时使用了某个gpu，加载的时候自动会经过这个gpu，但是这时候如果这个gpu内存不够的话就会报错
解决方法
在load_state_dict加一个变量

model.load_state_dict(torch.load("model/BERT_model.h5", map_location='cpu'))

这样就不会报CUDA内存不够的bug了，问题解决。
（ps：大家如果没有跑模型的时候记得把内存释放掉）