feat: 양자화 설정에 따라 device_map 자동 전환 기능 추가
This commit is contained in:
@@ -35,7 +35,7 @@ def download_model(config_path='./config.json'):
|
|||||||
model_name,
|
model_name,
|
||||||
cache_dir=local_path,
|
cache_dir=local_path,
|
||||||
quantization_config=quantization_config,
|
quantization_config=quantization_config,
|
||||||
device_map="cpu", # 다운로드 시 CPU에 로드하여 메모리 절약
|
device_map="auto" if quantization_config else "cpu", # 양자화 시 auto, 아니면 cpu
|
||||||
torch_dtype="auto"
|
torch_dtype="auto"
|
||||||
)
|
)
|
||||||
tokenizer = AutoTokenizer.from_pretrained(
|
tokenizer = AutoTokenizer.from_pretrained(
|
||||||
|
|||||||
Reference in New Issue
Block a user