dataset/downtask_segmentation/reverse.py

import numpy as np
import torch
import sys

file = sys.argv[1]

try:
    model = torch.load(file, map_location='cpu')['state_dict']
    torch.save(model, file)
    print('this file has been reversed.')
except KeyError:
    print('this file need not to reverse.')