from typing import Iterable
import torch


def _get_input_sizes(iterable: Iterable):
    """
    Gets the sizes of all torch tensors in an iterable. If an element
    of the iterable is a non-torch tensor iterable, it recurses into
    that iterable to continue calculating sizes. Any non-iterable is
    given a size of None. The output consists of nested lists with the
    same nesting structure as the input iterables.
    """
    out_list = []
    for i in iterable:
        if isinstance(i, torch.Tensor):
            out_list.append(list(i.size()))
        elif isinstance(i, Iterable):
            out_list.append(_get_input_sizes(i))
        else:
            out_list.append(None)
    return out_list