When I load a huge model like T5 xxl pretrained using device_map set to auto, and torch_dtype set to float16, it always insists on including my CPU when I have enough GPU ram (48 GB) how do I constrain accelerate to use only my GPUs? I tried setting the device_map manually with the layers spread over the GPUs but it gave an error.
2 Likes