该错误有可能是由于GPU不足导致的
有两种解决方法:

方法一：指定device

在指定device时，没有指定具体的卡
只用了如下代码

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

默认使用了index=0的卡，但是通过查看GPU使用情况发现
0号卡的占用率已经100%了，

所以只能换成其他卡跑
可以通过指定卡来限制，可以看到2号卡没有占用，所以我们可以指定2号卡
代码如下

device = torch.device("cuda:2" if torch.cuda.is_available() else "cpu")

方法二：调小batch size

这里举的例子比较特殊，正好GPU占用了100%，如果没有占用100%
我们也可以通过调小batch size来解决该问题。