File size: 263 Bytes
f7009b3
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
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.')