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

训练arcface时报这个错误。

原来是因为类别数没有设置正确。