mirror of
https://github.com/huggingface/text-generation-inference.git
synced 2025-04-22 23:42:06 +00:00
Update rw.py
This commit is contained in:
parent
198e6179ef
commit
377c01e21e
@ -39,7 +39,7 @@ class RW(CausalLM):
|
|||||||
device_map="auto"
|
device_map="auto"
|
||||||
if torch.cuda.is_available() and torch.cuda.device_count() > 1
|
if torch.cuda.is_available() and torch.cuda.device_count() > 1
|
||||||
else None,
|
else None,
|
||||||
load_in_8bit=quantize == "bitsandbytes",
|
load_in_4bit=quantize == "bitsandbytes",
|
||||||
trust_remote_code=trust_remote_code,
|
trust_remote_code=trust_remote_code,
|
||||||
)
|
)
|
||||||
if torch.cuda.is_available() and torch.cuda.device_count() == 1:
|
if torch.cuda.is_available() and torch.cuda.device_count() == 1:
|
||||||
|
Loading…
Reference in New Issue
Block a user