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

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

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

在这里插入图片描述

在这里插入图片描述

猜你喜欢

转载自blog.csdn.net/q506610466/article/details/127195815