I defined the function in my first post.
It does work for modules, such that you can call model.apply(weights_init).
In the method you could use conditions to initialize different module classes with different init functions.