lucas-ventura commited on
Commit
53d8b6c
·
1 Parent(s): 04f544e

Try with cublas

Browse files
Files changed (2) hide show
  1. app.py +9 -0
  2. requirements.txt +2 -0
app.py CHANGED
@@ -1,4 +1,6 @@
1
  import os
 
 
2
  import tempfile
3
  from pathlib import Path
4
  import subprocess
@@ -31,6 +33,13 @@ inference_model = None
31
 
32
  LLAMA_CKPT_PATH = "meta-llama/Meta-Llama-3.1-8B-Instruct"
33
 
 
 
 
 
 
 
 
34
  def install_cudnn():
35
  """Installs specific versions of libcudnn and configures torch for TF32."""
36
 
 
1
  import os
2
+ import nvidia.cublas.lib
3
+ import nvidia.cudnn.lib
4
  import tempfile
5
  from pathlib import Path
6
  import subprocess
 
33
 
34
  LLAMA_CKPT_PATH = "meta-llama/Meta-Llama-3.1-8B-Instruct"
35
 
36
+
37
+ cublas_path = os.path.dirname(nvidia.cublas.lib.__file__)
38
+ cudnn_path = os.path.dirname(nvidia.cudnn.lib.__file__)
39
+
40
+ ld_library_path = f"{cublas_path}:{cudnn_path}"
41
+ os.environ["LD_LIBRARY_PATH"] = ld_library_path
42
+
43
  def install_cudnn():
44
  """Installs specific versions of libcudnn and configures torch for TF32."""
45
 
requirements.txt CHANGED
@@ -1,5 +1,7 @@
1
  torch==2.5.1+cu121
2
  torchaudio==2.5.1+cu121
 
 
3
  faster-whisper
4
  pyannote.audio
5
  torchmetrics
 
1
  torch==2.5.1+cu121
2
  torchaudio==2.5.1+cu121
3
+ nvidia-cublas-cu11
4
+ nvidia-cudnn-cu11
5
  faster-whisper
6
  pyannote.audio
7
  torchmetrics