traindata = datasets.ImageFolder(data_dir + '/train', transforms.ToTensor())
image_means = torch.stack([t.mean(1).mean(1) for t, c in traindata])
image_means.mean(0)