import torch
from torchvision import d
from torch import nn
from torch.utils.data import dataset, dataloader
from torch import optim
train_loader = torch.utils.data.DataLoader(
torchvision.datasets.MNIST(’/files/’, train=True, download=True,
                         transform=torchvision.transforms.Compose([

                           torchvision.transforms.ToTensor(),

                           torchvision.transforms.Normalize(

                             (0.1307,), (0.3081,))

                         ])),

batch_size=64, shuffle=True)
val_loader = torch.utils.data.DataLoader(
torchvision.datasets.MNIST(’/files/’, train=False, download=True,
                         transform=torchvision.transforms.Compose([

                           torchvision.transforms.ToTensor(),

                           torchvision.transforms.Normalize(

                             (0.1307,), (0.3081,))

                         ])),

batch_size=1000, shuffle=True)

model
class CNN(nn.Module):
def init(self):
super(CNN, self).__init__()

self.conv = nn.Sequential(

    nn.Conv2d(1, 10, kernel_size=5), # non-RGB image hence input_channel = 1, out_channels is the number of feature maps,

    nn.ReLU(),

    nn.MaxPool2d(2),

    nn.Conv2d(10, 20, kernel_size=5), 

    nn.ReLU(),

    nn.MaxPool2d(2)

    )



self.flatten = nn.Flatten()

self.dense = nn.Sequential(

    nn.Linear(320, 50),

    nn.BatchNorm2d(50), 

    nn.Linear(50, 10)

    )

def forward(self, x):
x = self.conv(x)

x = self.flatten(x)

x = self.dense(x)

return x

model = CNN()
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())
epochs = 20
train_loss = []
val_loss = []
for t in range(epochs):
for feature, label in train_loader:
# forward

pred = model(feature)

# compute the objective function

output = loss_fn(pred, label)

# clean the gradients

model.zero_grad()

# backward propogation

output.backward()

# update

optimizer.step()

train_loss.append(output.item())

print(f’Epoch {t + 1}, train loss: {output}’)
for feature2, label2 in val_loader:
with torch.no_grad():

# forward

  pred2 = model(feature2)

# compute the objective function

output2 = loss_fn(pred2, label2)

val_loss.append(output2.item())

print(f'Epoch {t + 1}, train loss: {output}')

ValueError: expected 4D input (got 2D input)