chrisliu298

import collections
import logging

import torch
from transformers import AutoConfig, GPT2LMHeadModel, GPT2Tokenizer, set_seed

start = "<|startoftext|> "
sep = " <|sep|>"

chrisliu298

import numpy as np
from torchvision import datasets

train_dataset = datasets.CIFAR10("/tmp/data", train=True, download=True)
train_images = [np.array(train_dataset[i][0]) for i in range(len(train_dataset))]
train_images = np.stack(train_images)

means = (
    np.array(
        [