fix bf16
Browse files
app.py
CHANGED
|
@@ -19,7 +19,7 @@ from rwkv.model import RWKV
|
|
| 19 |
model_path = hf_hub_download(repo_id="Synthia/ChatGalRWKV", filename="rwkv-chatgal-v1-3B-ctx4096-epoch2.pth")
|
| 20 |
if 'ON_COLAB' in os.environ and os.environ['ON_COLAB'] == '1':
|
| 21 |
os.environ["RWKV_JIT_ON"] = '0'
|
| 22 |
-
os.environ["RWKV_CUDA_ON"] = '
|
| 23 |
model = RWKV(model=model_path, strategy='cuda bf16')
|
| 24 |
else:
|
| 25 |
model = RWKV(model=model_path, strategy='cpu bf16')
|
|
|
|
| 19 |
model_path = hf_hub_download(repo_id="Synthia/ChatGalRWKV", filename="rwkv-chatgal-v1-3B-ctx4096-epoch2.pth")
|
| 20 |
if 'ON_COLAB' in os.environ and os.environ['ON_COLAB'] == '1':
|
| 21 |
os.environ["RWKV_JIT_ON"] = '0'
|
| 22 |
+
os.environ["RWKV_CUDA_ON"] = '0' # if '1' then use CUDA kernel for seq mode (much faster)
|
| 23 |
model = RWKV(model=model_path, strategy='cuda bf16')
|
| 24 |
else:
|
| 25 |
model = RWKV(model=model_path, strategy='cpu bf16')
|