Save tensors to pickle file

save_tensors_to_pickle.py

def save_tensors_to_pickle(output_file: str, to_numpy: bool, **kwargs) -> None:
    """Save tensors to pickle file.
    
    Example usage:
    ```
    save_tensors_to_pickle(
        "train_output.pkl",
        to_numpy=False,
        inputs=inputs,
        logits=logits,
    )
    ```

    Args:
        output_file (str): output file path.
        to_numpy (bool): whether to convert tensors to numpy.
        **kwargs: tensors to save.
    """
    import pickle

    data = {}
    for key, value in kwargs.items():
        try:
            data[key] = value.cpu().detach()
            if to_numpy:
                data[key] = data[key].numpy()
        except Exception:
            pass

    with open(output_file, "wb+") as f:
        pickle.dump(data, f)