if type(module) in layers:

return {name: module}

res = {}

for name1, child in module.named_children():

res.update(find_layers(

child, layers=layers, name=name + '.' + name1 if name != '' else name1

))