One thing I would worry about is, if you are planning on using  (Distributed)DataParallel  at some point.
I’ve seen a lot of errors trying to just wrap the model into  DDP , if there were a lot of these “convenience functions”.

That is precisely what I’m dealing with at the moment indeed: the main module is wrapped in DP and I’ve had to disable that for now. Hope I’ll find a nice way to come around it