问题原因

报错信息：
OutOfMemoryError: CUDA out of memory. Tried to allocate 128.00 MiB (GPU 0; 6.00 GiB total capacity; 5.19 GiB already allocated; 0 bytes free; 5.21 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
原因：
cmd终端运行nvidia-smi命令，可以查看到CUDA已经被快被占满。
这种情况可能是之前的训练占用了显存，没被释放。

解决方法

使用torch.cuda.empty_cache()方法来释放PyTorch的缓存，以释放一部分被PyTorch占用的GPU内存。

import torch
torch.cuda.empty_cache()