Problems with models that don't have the parallelize() function

Hey,
First of all thank you for this amazing repo!
I am trying to employ this repo with a model that is does not have the parallelize() function (led - the longformer encoder-decoder).
Now - from what I have observed - such models are simply wrapped in a DataParallel decorator.
The problem is this causes many bugs that stem from the lack of parallelize function.
For example, many time the `get_policy_first_device` function is called, which searches for the `first_device` parameter in the models, which is inserted when `parallelize` is called (and there are many other issues).
I did notice a similar issue has already been addressed and so I was wondering if there are plans to properly treat such models.
Thanks!