For a single layer, this approach would work.
However, note that model.children() returns the child modules in the order they were initialized in the __init__, which might not be the same order they are called in the forward.