训练神经网络电脑自动重启

最近终于解决了一个困扰自己很长时间的问题。我在使用yolact的官方代码训练时，训练一个iteration后，电脑总会自动重启。
首先确定程序肯定没有问题，因为其他人利用此程序都是可以训练的。所以我最先想到的是是不是因为是多GPU训练的问题，后来测试之后并不是这个问题。然后将batch_size改小，还是会自动重启。今天突然想到可能是显存不够，所以尝试将图片的最大尺寸改小，由550改为300，特别神奇的是，可以运行了，不会自动重启了。
我可太开心了！！！！！！！
所以我觉得这样的问题应该是显存不足，可以尝试改小batch_size和image_size的方法

只是本人的初步尝试，如有错误欢迎指正。