Fix error in publish_model.py for pt>=1.6
@@ -20,6 +20,9 @@ def process_checkpoint(in_file, out_file):
del checkpoint['optimizer']
# if it is necessary to remove some sensitive data in checkpoint['meta'],
# add the code here.
+ if torch.__version__ >= '1.6':
+ torch.save(checkpoint, out_file, _use_new_zipfile_serialization=False)
+ else:
torch.save(checkpoint, out_file)
sha = subprocess.check_output(['sha256sum', out_file]).decode()
final_file = out_file.rstrip('.pth') + f'-{sha[:8]}.pth'
