Hi, I have the below CNN code, but I get an error when computing the Cross Entropy loss. Seems the shape of my y_hat is different from y_train, so how do I make the dimensions match?

X_train: torch.Size([1, 3, 708, 256])

Y_train: torch.Size([1, 708, 4])

```
class CNN(torch.nn.Module):
def __init__(self):
super(CNN,self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(3, 24, kernel_size=5, stride=1),
nn.Sigmoid(),
nn.MaxPool2d(kernel_size=2, stride=0))
self.conv2 = nn.Sequential(
nn.Conv2d(24, 12, kernel_size=5, stride=1),
nn.Sigmoid(),
nn.MaxPool2d(kernel_size=2, stride = 0))
#self.drop_out = nn.Dropout()
self.fc1 = nn.Linear(61, 12)
self.fc2 = nn.Linear(12, 4)
def forward(self, x):
out = self.conv1(x)
out = self.conv2(out)
#out = out.view(out.size(0), -1)
#out = out.reshape(out.size(0), -1)
#out = self.drop_out(out)
out = self.fc1(out)
out = self.fc2(out)
return out
```

```
# defining the model
model = CNN()
model = model.float()
# defining the optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=0.07)
# defining the loss function
criterion = CrossEntropyLoss()
# checking if GPU is available
if torch.cuda.is_available():
model = model.cuda()
criterion = criterion.cuda()
print(model)
```

```
def train(epoch):
model.train()
tr_loss = 0
# getting the training set
x_train, ytr= Variable(X_train), Variable(y_train)
ytrain = ytr.to(dtype=torch.int64)
# getting the validation set
#x_val, y_val = Variable(val_x), Variable(val_y)
# converting the data into GPU format
print('ytrain_shape',ytrain.shape)
# clearing the Gradients of the model parameters
optimizer.zero_grad()
# prediction for training and validation set
output_train = model(x_train.float())
# output_train = model(x_train)
output_train=output_train.squeeze(1)
print('Outputtrain shape',output_train.shape)
# computing the training and validation loss
loss_train = criterion(output_train, ytrain)
train_losses.append(loss_train)
# computing the updated weights of all the model parameters
loss_train.backward()
optimizer.step()
tr_loss = loss_train.item()
```

This part returns the error:

```
# defining the number of epochs
n_epochs = 5
# empty list to store training losses
train_losses = []
# empty list to store validation losses
val_losses = []
# training the model
for epoch in range(n_epochs):
train(epoch)
```