torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 128.00 MiB. GPU

这个错误提示表明在使用CUDA运算时，GPU的显存不足，无法分配128MB的显存。通常情况下，我们可以通过以下几种方式来解决这个问题：

1. 减少批次大小：可以减少每次训练的批次大小来降低显存使用量。

2. 减少模型参数：可以尝试减少模型的参数量或者使用更小的模型。

3. 释放不必要的显存：可以在每个batch结束后手动释放不必要的显存，比如清空一些变量或者删除一些无用的计算图节点。

4. 使用更高内存的GPU：如果你的GPU显存过小，可以尝试使用显存更大的GPU来避免内存不足。