from pathlib import Path
import torch
torch.manual_seed(42)
ROOT=Path(__file__).resolve().parents[1]
out=ROOT/'outputs'; out.mkdir(exist_ok=True)
features=['x1','x2','x3']
X=torch.tensor([[1.,2.,3.],[2.,4.,6.],[3.,6.,9.]])
mean=X.mean(dim=0); scale=X.std(dim=0,unbiased=False).clamp_min(1e-6)
model=torch.nn.Sequential(torch.nn.Linear(3,4),torch.nn.ReLU(),torch.nn.Linear(4,1))
model.eval()
with torch.inference_mode(): before=model((X-mean)/scale)
checkpoint={'model_state':model.state_dict(),'mean':mean,'scale':scale,'features':features,'torch_version':str(torch.__version__)}
path=out/'torch-save.pt'; torch.save(checkpoint,path)
loaded=torch.load(path,map_location='cpu',weights_only=True)
assert loaded['features']==features
restored=torch.nn.Sequential(torch.nn.Linear(3,4),torch.nn.ReLU(),torch.nn.Linear(4,1))
restored.load_state_dict(loaded['model_state']); restored.eval()
with torch.inference_mode(): after=restored((X-loaded['mean'])/loaded['scale'])
torch.testing.assert_close(before,after)
assert torch.isfinite(after).all()
print('checkpoint=',path.name,'prediction_shape=',tuple(after.shape),'roundtrip=PASS')
print('predictions=',after.flatten().tolist())