Fix balanced and auto device_map
# What does this PR do?

In #22095 some of the arguments passed to `infer_auto_device_map` were grouped in kwargs. The problem is that one of those (`max_memory`) is not updated anymore after being changed (when device_map is `"auto"`, `"balanced"` or `"balanced_low_0"`). This PR fixes that.

Note: this is a regression so this will need to go in a patch.
_The documentation is not available anymore as the PR was closed or merged._