from typing import Union
import torch


def check_data_types(data) -> Union[list, torch.Tensor]:
    """
    Utility function to check if data passed is all on tensor format. Can handle iterables or single structures

    :param data: Data to be checked
    :return: list of tensors
    """

    if torch.is_tensor(data):
        return data

    else:
        try:
            return torch.tensor(data, dtype=torch.float)
        except ValueError:
            return [torch.tensor(x, dtype=torch.float) for x in data]