Solved by ptrblck in post #2 
              
              
                Yes, manually unwrapping the data should work. Something like this might help: 
# flatten
gg = []
for name, weight in network.named_parameters():
    gg.extend(weight.cpu().detach().numpy().flatten())
    
# restore
sd = {}
i = 0
for name, weight in network.named_parameters():
    param_len = weight…