I changed the normalize code from
def normalize(t, mean, std, dataset):

    if dataset == 'mnist':
        t[:, 0, :] = (t[:, 0, :] - mean[0]) / std[0]
    if dataset == 'cifar10':
        t[:, 0, :, :] = (t[:, 0, :, :].detach() - mean[0]) / std[0]
        t[:, 1, :, :] = (t[:, 1, :, :].detach() - mean[1]) / std[1]
        t[:, 2, :, :] = (t[:, 2, :, :].detach() - mean[2]) / std[2]

    return t

to this
def normalize(t, mean, std, dataset):
    return t

and It is working, so I assuming that I should rewrite the code inside the normalize function differently.