PS: I noticed after posting the above that this line:
layer_new = layer_type_new(module.in_channels, module.out_channels, module.kernel_size, module.stride,
module.padding, module.dilation, module.groups,
module.bias is not None, module.padding_mode) c
will only work if the layer to replace is nn.Conv2d
and the new layer has similar parameters to pass to its constructor. If someone has a more generic way to handle this that will be great to share.