Layer Groups not saved when exporting model

I just tested this. When I export my model, and reload it back, the layer groups are not saved.

Previously I trained a U-Net, and there was 3 layer groups, which allows me to freeze/unfreeze the network. Now when I reload the exported model, everything is lumped into 1 layer group.

Does anyone know a workaround for this for saving/restoring the original layer group configuration?


have the save problem with CNN restnet34. Interested to find out the solution

I have found the following solution: learn.split(lambda m: (m[0][6], m[1])). This returns 3 layer groups for loaded models