Merge pull request #206 from ByeongkiJeong/patch-1

Update inference_on_multiple_gpus.md
This commit is contained in:
Cui Junbo
2024-06-04 22:27:15 +08:00
committed by GitHub

View File

@@ -43,7 +43,7 @@ gpu_device_ids = [0, 1] # Define which gpu to use (now we have two GPUs, each ha
no_split_module_classes = ["LlamaDecoderLayer"]
max_memory = {
device_id: memory for device_id in gpu_device_ids
device_id: max_memory_each_gpu for device_id in gpu_device_ids
}
config = AutoConfig.from_pretrained(