CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 2.00 GiB total capacity; 1.04 GiB already allocated; 0 bytes free; 1.08 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

这个错误表明你的GPU内存不足，PyTorch试图分配20.00 MiB的内存却失败了。已经有1.04 GiB的内存被分配了，没有剩余内存可用。你可以尝试使用 max_split_size_mb 参数设置最大拆分大小以避免内存碎片化。PYTORCH_CUDA_ALLOC_CONF提供了内存管理方面的文档，你可以参考它来解决这个问题。此外，你也可以尝试减小模型的批次大小或者使用更小的模型来减少内存占用。