import torch, os

pred_path = "/mnt/shared-storage-user/caipengxiang/workspace/ChemBOMAS/train_regression/data4regression/suzuki_60/cluster_pretrain_yields.pt"
basename = os.path.basename(pred_path).split(".pt")[0] + "_merged.pt"
data = torch.load(pred_path, map_location='cpu')

split_path = "split_idx.pt"
split_data = torch.load(split_path)

data['train_idx'] = split_data['train']
data['val_idx'] = split_data['test']
torch.save(data, basename)
