I set batch_size to 1 and it works.
When I set batch_size to 2, and the same error is reported again:
“RuntimeError: CUDA out of memory. Tried to allocate 30.00 MiB (GPU 0; 2.00 GiB total capacity; 784.67 MiB already allocated; 551.43 MiB free; 31.33 MiB cached)”.
I use nvidia-smi and only this process exists.
Do you have other suggestions?
Thanks!