Fix : get_balanced_memory when using multi gpus with small models or quantized models with a large vocabulary

Fix : get_balanced_memory when using multi gpus with small models or quantized models with a large vocabulary #5398