How to pass device_id to overriden functions?

I want to select the gpu that my custom model is being loaded into. Can't find a way to do this, as when we do self.load_fn = load_fn we lose the reference to self inside load_fn

Start by tracing the custom model loading path around `self.load_fn = load_fn` and inspect how overridden functions are invoked. Determine how `device_id` could remain available to the custom loader, and consider the work done when a custom model can select the GPU it is loaded into.

Written by the indexing model from the issue text.