Maybe you can refer to the new features of PyTorch, torch.package

https://docs.pytorch.org/docs/stable/package.html

import torch.package

# save
model = YourModel()
pkg = torch.package.PackageExporter("model_package.pt")
pkg.save_pickle("model", "model.pkl", model)

import torch.package
import sys
import importlib.util

# load
imp = torch.package.PackageImporter("model_package.pt")
model = imp.load_pickle("model", "model.pkl")