Ошибка CUDA: CUBLAS_STATUS_ALLOC_FAILED при вызове `cublasCreate (handle)`

Question

Я получил следующую ошибку, когда запустил свою модель глубокого обучения pytorch в colab

/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py in linear(input, weight, bias)
   1370         ret = torch.addmm(bias, input, weight.t())
   1371     else:
-> 1372         output = input.matmul(weight.t())
   1373         if bias is not None:
   1374             output += bias

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

Я даже уменьшил размер пакета со 128 до 64, то есть уменьшил до половины, но все же я получил эту ошибку. Ранее я запускал тот же код с размером пакета 128, но не получал никакой ошибки, подобной этой.