import torch
chunk0 = torch.load('./checkpoints_00335001_model_chunk0.pth')
chunk1 = torch.load('./checkpoints_00335001_model_chunk1.pth')
others = torch.load('./checkpoints_00335001_others.pth')

model = {**chunk0, **chunk1}
res = {'model': model, **others}
torch.save(res, './checkpoints_00335001.pth')