Calling infer_auto_device_map() with max_memory=None throws an error in version 1.2.0

[chmeyers]

System Info

accelerate==1.2.0

Reproduction

Bug is from this commit:
d7b1b36

max_memory initialization was moved into _init_infer_auto_device_map, which does not return the max_memory value.

So if max_memory=None is passed to infer_auto_device_map (the default value), then it will still be None at line 1415:

accelerate/src/accelerate/utils/modeling.py

Line 1415 in cb8b7c6

current_max_size = max_memory[device] if device != "disk" else None

Leading to error: TypeError: 'NoneType' object is not subscriptable

Expected behavior

max_memory=None when passed to infer_auto_device_map does not throw an error.

[chmeyers]

@Nech-C

[Nech-C]

Sorry for the oversight. I will fix it ASAP. Thanks for pointing it out!