Hello, below is my code for MPL image classification, I am running the code and I forgot to define the name, "device." (Highlighted in bold) Can someone look over it to see where it would be best to include the definition? Thank You.

import torchvision.datasets as datasets

import torch.optim as optim

import torch.utils.data as data

import torch.nn as nn

from torchvision import transforms

train_transform = transforms.Compose([

transforms.RandomCrop(32, padding=4),

transforms.RandomHorizontalFlip(),

transforms.ToTensor(),

])

test_transform = transforms.Compose([

transforms.ToTensor(),

])

train_data = datasets.CIFAR100(root='data', train=True, transform=train_transform, download=True)

test_data = datasets.CIFAR100(root='data', train=False, transform=test_transform, download=True)

class MLP(nn.Module):

def __init__(self, input_size, hidden_size, num_classes):

super(MLP, self).__init__()

self.fc1 = nn.Linear(input_size, hidden_size)

self.relu1 = nn.ReLU()

self.fc2 = nn.Linear(hidden_size, hidden_size)

self.relu2 = nn.ReLU()

self.fc3 = nn.Linear(hidden_size, num_classes)

def forward(self, x):

out = self.fc1(x)

out = self.relu1(out)

out = self.fc2(out)

out = self.relu2(out)

out = self.fc3(out)

return out

input_size = 32 * 32 * 3

hidden_size = 512

num_classes = 100

learning_rate = 0.001

batch_size = 128

num_epochs = 10

model = MLP(input_size, hidden_size, num_classes)

train_loader = data.DataLoader(train_data, batch_size=batch_size, shuffle=True)

test_loader = data.DataLoader(test_data, batch_size=batch_size, shuffle=False)

criterion = nn.CrossEntropyLoss()

optimizer = optim.Adam(model.parameters(), lr=learning_rate)

for epoch in range(num_epochs):

for i, (images, labels) in enumerate(train_loader):

images = images.reshape(-1, input_size).to(device)

labels = labels.to(device)

optimizer.zero_grad()

outputs = model(images)

loss = criterion(outputs, labels)

loss.backward()