Change CSE to use torch.save instead of pickle.dump
Summary: Change CSE to use torch.save instead of pickle.dump to be compatible with torch.load
LaMa Project: L1124100
@@ -5,7 +5,6 @@ import argparse
import glob
import logging
import os
-import pickle
import sys
from typing import Any, ClassVar, Dict, List
import torch
@@ -190,7 +189,7 @@ class DumpAction(InferenceAction):
if len(out_dir) > 0 and not os.path.exists(out_dir):
os.makedirs(out_dir)
with open(out_fname, "wb") as hFile:
- pickle.dump(context["results"], hFile)
+ torch.save(context["results"], hFile)
logger.info(f"Output saved to {out_fname}")
