Would it make sense to iterate through the model using
for name, layer in module.named_modules():
    pass

in combination with getattr() and setattr() or what would be the PyTorch way of replacing layers in arbitrary networks?