def _find_tensors(
    obj: Union[Tensor, list, tuple, dict, Any],
) -> Union[list[Tensor], itertools.chain]:  # pragma: no-cover
    """Recursively find all tensors contained in the specified object."""
    if isinstance(obj, Tensor):
        return [obj]
    if isinstance(obj, (list, tuple)):
        return itertools.chain(*map(_find_tensors, obj))
    if isinstance(obj, dict):
        return itertools.chain(*map(_find_tensors, obj.values()))
    return []
Recursively find all tensors contained in the specified object.
src/lightning/pytorch/overrides/distributed.py