ifsheldon/stannum

PyTorch checkpointing with `Tin`

ifsheldon opened this issue · 0 comments

Now although Tin is subclass of torch.nn.Module, we cannot actually save parameters and "model structures" in Tin because parameters are values in Taichi fields and "model structures" are actually Taichi kernels and Taichi functions.

To totally serialize a Tin, we need to serialize parameters, which may be achieved by a proxy torch.nn.Parameter (see issue #3 ), but we also need to serialize Taichi kernels and Taichi functions.

If anyone comes up with a better solution, discussions and PRs are always welcomed.