Purpose

• Enable util that may be useful for dealing with offloading of modules which are no leaf modules. For example, if we want to attach parameters to an attention module for attention quantization, we'll need to know the offload device of the attention module (which is not a leaf module)

Changes

• Generalize get_offloaded_device to support nested modules

Testing

• Added additional tests, previous tests pass and previous behavior is preserved