from typing import Union
import torch
import pathlib


def export_model(model: Union[torch.nn.Module, torch.jit.ScriptModule],
                 path: Union[str, pathlib.Path],
                 input_shape: tuple = (1, 3, 64, 64),
                 use_script_module: bool = True
                 ) -> Union[str, pathlib.Path]:
    """Exports model.

    Args:
        model (torch.nn.Module or torch.jit.ScriptModule): Saved model.
        path (str or pathlib.Path): Path to file.
        input_shape (tuple, optional): Tuple of input data shape.
        use_script_module (bool, optional): Boolean flag for using script
            module.

    Returns:
        path (str or pathlib.Path): Path to saved file.
    """

    model = model.cpu().eval()
    if isinstance(model, torch.jit.ScriptModule):
        assert use_script_module, \
            "Provided model is a ScriptModule, set use_script_module to True."
    if use_script_module:
        if not isinstance(model, torch.jit.ScriptModule):
            assert input_shape is not None
            traced_model = torch.jit.trace(model, torch.zeros(*input_shape))
        else:
            traced_model = model
        torch.jit.save(traced_model, path)
    else:
        torch.save(model, path)  # saves model as a nn.Module
    return path