Partition data 1

import math
import numpy as np
import torch
from torch import nn
import matplotlib.pyplot as plt
import torchvision.datasets as datasets
from torch.utils import data
from torchvision import transforms
from torch.utils.data import Dataset
import torchvision


train_data=datasets.CIFAR10(root="../data", train=True, transform=transforms.ToTensor(), download=True)
test_data=datasets.CIFAR10(root="../data", train=False, transform=transforms.ToTensor(), download=True)
t0=[ i for i, x in enumerate(train_data.targets) if x == 0]
t1=[ i for i, x in enumerate(train_data.targets) if x == 1]
t2=[ i for i, x in enumerate(train_data.targets) if x == 2]
t3=[ i for i, x in enumerate(train_data.targets) if x == 3]
t4=[ i for i, x in enumerate(train_data.targets) if x == 4]
t5=[ i for i, x in enumerate(train_data.targets) if x == 5]
t6=[ i for i, x in enumerate(train_data.targets) if x == 6]
t7=[ i for i, x in enumerate(train_data.targets) if x == 7]
t8=[ i for i, x in enumerate(train_data.targets) if x == 8]
t9=[ i for i, x in enumerate(train_data.targets) if x == 9]


label0=[]
for i in t0:
    label0.append(train_data.targets[i])
label1=[]
for i in t1:
    label1.append(train_data.targets[i])
label2=[]
for i in t2:
    label2.append(train_data.targets[i])
label3=[]
for i in t3:
    label3.append(train_data.targets[i])
label4=[]
for i in t4:
    label4.append(train_data.targets[i])
label5=[]
for i in t5:
    label5.append(train_data.targets[i])
label6=[]
for i in t6:
    label6.append(train_data.targets[i])
label7=[]
for i in t7:
    label7.append(train_data.targets[i])
label8=[]
for i in t8:
    label8.append(train_data.targets[i])
label9=[]
for i in t9:
    label9.append(train_data.targets[i])

data0=np.array([])
for i in t0:
    data0=np.append(data0,train_data.data[i])
data1=np.array([])
for i in t1:
    data1=np.append(data1,train_data.data[i])
data2=np.array([])
for i in t2:
    data2=np.append(data2,train_data.data[i])
data3=np.array([])
for i in t3:
    data3=np.append(data3,train_data.data[i])
data4=np.array([])
for i in t4:
    data4=np.append(data4,train_data.data[i])
data5=np.array([])
for i in t5:
    data5=np.append(data5,train_data.data[i])
data6=np.array([])
for i in t6:
    data6=np.append(data6,train_data.data[i])
data7=np.array([])
for i in t7:
    data7=np.append(data7,train_data.data[i])
data8=np.array([])
for i in t8:
    data8=np.append(data8,train_data.data[i])
data9=np.array([])
for i in t9:
    data9=np.append(data9,train_data.data[i])

data0=data0.reshape(len(label0),32,32,3)
data1=data1.reshape(len(label1),32,32,3)
data2=data2.reshape(len(label2),32,32,3)
data3=data3.reshape(len(label3),32,32,3)
data4=data4.reshape(len(label4),32,32,3)
data5=data5.reshape(len(label5),32,32,3)
data6=data6.reshape(len(label6),32,32,3)
data7=data7.reshape(len(label7),32,32,3)
data8=data8.reshape(len(label8),32,32,3)
data9=data9.reshape(len(label9),32,32,3)

class DatasetXY(Dataset):
    def __init__(self, x, y):
        self._x = x
        self._y = y
        self._len = len(x)

    def __getitem__(self, item):  #  The value returned in each cycle 
        return self._x[item], self._y[item]

    def __len__(self):
        return self._len

dataset0= DatasetXY(data0,label0)
dataset1= DatasetXY(data1,label1)
dataset2= DatasetXY(data2,label2)
dataset3= DatasetXY(data3,label3)
dataset4= DatasetXY(data4,label4)
dataset5= DatasetXY(data5,label5)
dataset6= DatasetXY(data6,label6)
dataset7= DatasetXY(data7,label7)
dataset8= DatasetXY(data8,label8)
dataset9= DatasetXY(data9,label9)

train1_iter=data.DataLoader(dataset0,batch_size=32,num_workers=0)