dexifried commited on
Commit
da90885
·
1 Parent(s): 018a4d9

Ensure GPU detection with warmup at init

Browse files
Files changed (1) hide show
  1. app.py +5 -1
app.py CHANGED
@@ -37,7 +37,11 @@ SYSTEM = (
37
  @spaces.GPU(duration=60)
38
  def gpu_warmup() -> str:
39
  import torch
40
- return f"GPU ready: {torch.cuda.get_device_name(0)}"
 
 
 
 
41
 
42
  @spaces.GPU(duration=120)
43
  def _gpu_llm_chat(messages: list) -> str:
 
37
  @spaces.GPU(duration=60)
38
  def gpu_warmup() -> str:
39
  import torch
40
+ return f"GPU ready at init: {torch.cuda.is_available()}"
41
+
42
+ gpu_message = gpu_warmup() # Ensure detection during init
43
+ logger.info(gpu_message)
44
+
45
 
46
  @spaces.GPU(duration=120)
47
  def _gpu_llm_chat(messages: list) -> str: