how to save and load the checkpoint of splitNN?

[lepangdan]

Question

How to save and load the checkpoint of splitNN?

Further Information

Even though it is easy to save and load models in vanilla PyTorch, I encountered some trouble when I did the same thing in hook

Here is my code and corresponding result, I couldn't get the value of weights, instead of wrapped pointer Tensor. Please help

def state_dict(self):
        res = {}
        for model in self.models:
            res = {k: v for k, v in model.state_dict().items()}
        for k, v in res.items():
            print('k, v', k, v)
        return res

k, v final_layer.bias (Wrapper)>[PointerTensor | me:14397441124 -> bob:14831367252]::data
models_state_dict {'final_layer.weight': (Wrapper)>[PointerTensor | me:42574993002 -> bob:7259753333]::data, 'final_layer.bias': (Wrapper)>[PointerTensor | me:14397441124 -> bob:14831367252]::data}
k, v final_layer.weight (Wrapper)>[PointerTensor | me:42574993002 -> bob:7259753333]::data
k, v final_layer.bias (Wrapper)>[PointerTensor | me:14397441124 -> bob:14831367252]::data