The previous code can run normally. After the data set is expanded, the following errors are reported in the GPU program running the deep learning training model, but CUDA out of memory error is not prompted.

RuntimeError: CUDA error: CUBLAS_STATUS_ALLOC_FAILED when calling `cublasCreate(handle)`

[solution 1]

Change the program to run on the CPU and find that it can run normally, but the speed will be very slow and it will take a long time.

--device cpu

[solution 2]

Try to reduce the batch size used in the training model, and it can run normally.