Expected tensor for argument #1 'input' to have the same dimension

Using below code I create 10 instances of each training data of which each has 100 dimensions. Each of the 100 dimension contains 3 dimensions. Therefore it's shape is : (3, 100, 10). This emulates 10 instances of 100 pixels each with 3 channels to emulate an RGB value

I've set this model to just classify between 1 and 0.

When applying the softmax layer I receive the error :

RuntimeError: Expected tensor for argument #1 'input' to have the same dimension as tensor for 'result'; but 4 does not equal 3 (while checking arguments for cudnn_convolution)

I'm using 0.4.0 (checked using print(torch.__version__) ) How to correctly set the dimension for the softmax layer? As I think my dimensions are correct?

%reset -f

import os
import torch
from skimage import io, transform
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
import torch.utils.data as data_utils
import torchvision
import numpy as np
from sklearn.preprocessing import scale
import torch.nn.functional as F
import torch.nn as nn
import torch.nn.functional as F
from random import randint

batch_size_value = 10

train_dataset = []
mu, sigma = 0, 0.1 # mean and standard deviation
num_instances = 10

# Create 3000 instance and reshape to (3 , 100, 10) , this emulates 10 instances of 100 pixels 
# each with 3 channels to emulate an RGB value

for i in range(num_instances) :
    image = []
    image_x = np.random.normal(mu, sigma, 3000).reshape((3 , 100, 10))
    train_dataset.append(image_x)

mu, sigma = 100, 0.80 # mean and standard deviation
for i in range(num_instances) :
    image = []
    image_x = np.random.normal(mu, sigma, 3000).reshape((3 , 100, 10))
    train_dataset.append(image_x)

labels_1 = [1 for i in range(num_instances)]
labels_0 = [0 for i in range(num_instances)]

labels = labels_1 + labels_0

print(labels)

x2 = torch.tensor(train_dataset).float()
y2 = torch.tensor(labels).long()

my_train2 = data_utils.TensorDataset(x2, y2)
train_loader2 = data_utils.DataLoader(my_train2, batch_size=batch_size_value, shuffle=False)

# print(x2)

# Device configuration
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('device' , device)
# device = 'cpu'

# Hyper parameters
num_epochs = 10
num_classes = 2
batch_size = 5
learning_rate = 0.001

# Convolutional neural network (two convolutional layers)
class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, kernel_size=5) 
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
        self.fc1   = nn.Linear(864, 120) 
        self.fc2   = nn.Linear(120, 84)
        self.fc3   = nn.Linear(84, num_classes)

    def forward(self, x):
        out = F.relu(self.conv1(x))
        out = F.max_pool2d(out, 2)
        out = F.relu(self.conv2(out))
        out = F.max_pool2d(out, 2)
        out = out.view(out.size(0), -1)
        out = F.relu(self.fc1(out))
        out = F.relu(self.fc2(out))
        out = self.fc3(out)

model = ConvNet().to(device)

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# Train the model
total_step = len(train_loader2)
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader2):
        images = images.to(device)
        labels = labels.to(device)

        # Forward pass
        outputs = model(images)
#         print(images)

        loss = criterion(outputs, labels)

        # Backward and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (i % 10) == 0:
            print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' 
                   .format(epoch+1, num_epochs, i+1, total_step, loss.item()))

from Expected tensor for argument #1 'input' to have the same dimension