LutaoJiang commited on
Commit
695826d
Β·
1 Parent(s): 548d313
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -77,6 +77,7 @@ def check_gpu():
77
  print("Device count:", torch.cuda.device_count())
78
 
79
  # test nvdiffrast
 
80
  import nvdiffrast.torch as dr
81
  dr.RasterizeCudaContext(device="cuda:0")
82
  print("nvdiffrast initialized successfully")
@@ -177,6 +178,7 @@ def text_to_detailed(prompt, seed=None):
177
  "rm -rf /home/user/.cache/torch_extensions/py310_cu121/nvdiffrast_plugin/lock"
178
  )
179
  )
 
180
  if "CUDA_VISIBLE_DEVICES" in os.environ:
181
  del os.environ["CUDA_VISIBLE_DEVICES"]
182
  subprocess.run(['nvidia-smi'])
 
77
  print("Device count:", torch.cuda.device_count())
78
 
79
  # test nvdiffrast
80
+ os.environ["TORCH_CUDA_ARCH_LIST"] = "9.0"
81
  import nvdiffrast.torch as dr
82
  dr.RasterizeCudaContext(device="cuda:0")
83
  print("nvdiffrast initialized successfully")
 
178
  "rm -rf /home/user/.cache/torch_extensions/py310_cu121/nvdiffrast_plugin/lock"
179
  )
180
  )
181
+ os.environ["TORCH_CUDA_ARCH_LIST"] = "9.0"
182
  if "CUDA_VISIBLE_DEVICES" in os.environ:
183
  del os.environ["CUDA_VISIBLE_DEVICES"]
184
  subprocess.run(['nvidia-smi'])