def _convert_tensor_args(params: Union[List[TorchTensor], Dict[str, TorchTensor]]):
    """
    Helper function for the WithTensors APIs.
    """
    if isinstance(params, dict):
        result = {name: torch_to_ait_data(x) for name, x in params.items()}
    else:
        result = [torch_to_ait_data(x) for x in params]
    return result

    Helper function for the WithTensors APIs.
    
python/aitemplate/compiler/model.py