picocreator commited on
Commit
42e5f71
1 Parent(s): f787cc5

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -1
app.py CHANGED
@@ -24,11 +24,14 @@ except NVMLError as error:
24
 
25
 
26
  os.environ["RWKV_JIT_ON"] = '1'
27
- os.environ["RWKV_CUDA_ON"] = '1' # if '1' then use CUDA kernel for seq mode (much faster)
28
 
29
  # Model strat to use
30
  MODEL_STRAT="cpu bf16"
 
 
 
31
  if HAS_GPU == True :
 
32
  MODEL_STRAT = "cuda bf16"
33
 
34
  # Load the model accordingly
 
24
 
25
 
26
  os.environ["RWKV_JIT_ON"] = '1'
 
27
 
28
  # Model strat to use
29
  MODEL_STRAT="cpu bf16"
30
+ os.environ["RWKV_CUDA_ON"] = '0' # if '1' then use CUDA kernel for seq mode (much faster)
31
+
32
+ # Switch to GPU mode
33
  if HAS_GPU == True :
34
+ os.environ["RWKV_CUDA_ON"] = '1'
35
  MODEL_STRAT = "cuda bf16"
36
 
37
  # Load the model accordingly