abreza commited on
Commit
d6247cd
·
1 Parent(s): 9835f74
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -109,7 +109,8 @@ def find_cuda_home():
109
 
110
  @spaces.GPU
111
  def upsample_features(image, model_option):
112
- os.environ["CUDA_HOME"] = find_cuda_home()
 
113
  print(os.environ["CUDA_HOME"])
114
  os.environ['PATH'] = '/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin'
115
  os.environ['LD_LIBRARY_PATH'] = '/usr/local/nvidia/lib:/usr/local/nvidia/lib64'
 
109
 
110
  @spaces.GPU
111
  def upsample_features(image, model_option):
112
+ from torch.utils.cpp_extension import _find_cuda_home
113
+ os.environ["CUDA_HOME"] = _find_cuda_home()
114
  print(os.environ["CUDA_HOME"])
115
  os.environ['PATH'] = '/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin'
116
  os.environ['LD_LIBRARY_PATH'] = '/usr/local/nvidia/lib:/usr/local/nvidia/lib64'