from typing import Iterable
import torch
import itertools


def compute_mean_std(images: Iterable[torch.Tensor], nsample: int = -1):
    """
    Compute the mean and standard deviation for a data loader of image tensors.
    """
    if nsample == -1:
        images = images
    else:
        images = itertools.islice(images, nsample)
    X = torch.stack(list(images))
    X = torch.transpose(X, 0, 1).contiguous().view(3, -1)
    mean = torch.mean(X, dim=1)
    std = torch.std(X, dim=1)
    return mean.numpy(), std.numpy()