Using 2 GPUs out of 4

Beginners
1

Hi,

I’m doing inference on falcon 40B.
I have 4 GPU on my server.
when I do : model = AutoModelForCausalLM.from_pretrained(model_id, cache_dir=‘./workspace/’,
torch_dtype=torch.bfloat16, trust_remote_code=True, device_map=auto, offload_folder=“offload”)
it is using the 4 GPU.
When I do : model = AutoModelForCausalLM.from_pretrained(model_id, cache_dir=‘./workspace/’,
torch_dtype=torch.bfloat16, trust_remote_code=True, device_map={“” :0}, offload_folder=“offload”)
it is using only 1 GPU

Now I don’t now the syntax of device_map to tell him to use only 2 GPU (0 and 1).

thanks for your help