Intel Gpu support initialization (#4340)
This commit is contained in:
parent
317e2c857e
commit
778a010df8
14 changed files with 106 additions and 42 deletions
|
@ -9,6 +9,7 @@ from pathlib import Path
|
|||
|
||||
import numpy as np
|
||||
from tokenizers import Tokenizer
|
||||
from transformers import is_torch_xpu_available
|
||||
|
||||
import modules.shared as shared
|
||||
from modules.callbacks import Iteratorize
|
||||
|
@ -27,7 +28,7 @@ class RWKVModel:
|
|||
pass
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(self, path, dtype="fp16", device="cuda"):
|
||||
def from_pretrained(self, path, dtype="bf16" if is_torch_xpu_available() else "fp16", device="xpu" if is_torch_xpu_available() else "cuda"):
|
||||
tokenizer_path = Path(f"{path.parent}/20B_tokenizer.json")
|
||||
if shared.args.rwkv_strategy is None:
|
||||
model = RWKV(model=str(path), strategy=f'{device} {dtype}')
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue