#! /usr/bin/env python
# -*- coding: utf-8 -*-
'''
Created on 2023/4/21 11:28:41
@author: LiuKuan
@copyright: Apache License, Version 2.0
'''
from laok.third_lib import torch
#===============================================================================
r'''
'''
#===============================================================================
__all__ = ['load_torch_file', 'save_torch_file']

def load_torch_file(filename, map_location = None, pickle_module = None, *, weights_only = False, **pickle_load_args):
    return torch.load(f=filename, map_location=map_location,
               pickle_module=pickle_module, weights_only=weights_only, **pickle_load_args)

def save_torch_file(filename, obj, **kws):
    torch.save(obj=obj, f=filename, **kws)


def load_torch_jit(filename):
    ''' load trace model
    '''
    return torch.jit.load(filename)

def save_torch_jit(model, input, filename):
    if isinstance(model, torch.jit.TracedModule):
        traced_model = model
    else:
        try:
            mode = model.training
            model.eval()
            traced_model = torch.jit.trace(model, input)
        finally:
            model.train(mode)

    traced_model.save(filename)
    return traced_model

def save_torch_onnx(model, input, filename, **kws):
    torch.onnx.export(model, input, filename, **kws)