CUBLAS_STATUS_INVALID_VALUE cuda error when running on aws p2 instance

[reubenfirmin]

I installed whisper on a p2.xlarge running AWS's "AWS Deep Learning AMI GPU CUDA 11.4.3 (Ubuntu 20.04) 20220210" AMI. Running transcription using cuda, I get:

env/bin/whisper --language=English 4e3cc005-8573-4d5f-9175-625c0ae00bdd_upload_file__x0.mp3 --device=cuda --model=base
Traceback (most recent call last):
  File "env/bin/whisper", line 8, in <module>
    sys.exit(cli())
  File "/home/ubuntu/env/lib/python3.8/site-packages/whisper/transcribe.py", line 307, in cli
    result = transcribe(model, audio_path, temperature=temperature, **args)
  File "/home/ubuntu/env/lib/python3.8/site-packages/whisper/transcribe.py", line 181, in transcribe
    result: DecodingResult = decode_with_fallback(segment)
  File "/home/ubuntu/env/lib/python3.8/site-packages/whisper/transcribe.py", line 117, in decode_with_fallback
    decode_result = model.decode(segment, options)
  File "/home/ubuntu/env/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/ubuntu/env/lib/python3.8/site-packages/whisper/decoding.py", line 701, in decode
    result = DecodingTask(model, options).run(mel)
  File "/home/ubuntu/env/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/ubuntu/env/lib/python3.8/site-packages/whisper/decoding.py", line 617, in run
    audio_features: Tensor = self._get_audio_features(mel)  # encoder forward pass
  File "/home/ubuntu/env/lib/python3.8/site-packages/whisper/decoding.py", line 561, in _get_audio_features
    audio_features = self.model.encoder(mel)
  File "/home/ubuntu/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/env/lib/python3.8/site-packages/whisper/model.py", line 156, in forward
    x = block(x)
  File "/home/ubuntu/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/env/lib/python3.8/site-packages/whisper/model.py", line 124, in forward
    x = x + self.attn(self.attn_ln(x), mask=mask, kv_cache=kv_cache)
  File "/home/ubuntu/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/env/lib/python3.8/site-packages/whisper/model.py", line 78, in forward
    k = self.key(x if xa is None else xa)
  File "/home/ubuntu/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/env/lib/python3.8/site-packages/whisper/model.py", line 36, in forward
    return F.linear(
RuntimeError: CUDA error: CUBLAS_STATUS_INVALID_VALUE when calling `cublasSgemmEx( handle, opa, opb, m, n, k, &falpha, a, CUDA_R_16F, lda, b, CUDA_R_16F, ldb, &fbeta, c, CUDA_R_16F, ldc)`

[daxaxelrod]

Getting something similar with a T4 on GCP