Closed gsarti closed 7 months ago
This PR adds support for using a custom device_map in the loaded HF model, avoiding to cast it to a device in AttributionModel.setup() if the device map is specified, and dispatching it using the accelerate.dispatch_model method instead.
device_map
AttributionModel.setup()
accelerate.dispatch_model
Description
This PR adds support for using a custom
device_mapin the loaded HF model, avoiding to cast it to a device inAttributionModel.setup()if the device map is specified, and dispatching it using theaccelerate.dispatch_modelmethod instead.