修改huggingface 模型output_hidden_states=True时,CUDA out of memory问题。

在使用huggingface 提供的Trainer进行模型预测时,如果训练时output_hidden_states=True,则显存占用会无限增加,最终导致CUDA out of memory 内存溢出错误。
解决办法:

在模型的最终返回值处,把hidden_states 设置为None 即可,具体原因我也不清楚。

修改huggingface 模型output_hidden_states=True时,CUDA out of memory问题。_第1张图片

修改huggingface 模型output_hidden_states=True时,CUDA out of memory问题。_第2张图片

你可能感兴趣的:(深度学习,深度学习,人工智能)