I am newbie in pytorch, I am trying to create a convolutional neural network with my own data but I have a problem running the cross entropy loss function. I am not sure but maybe it is due to the structure of my data. Data train contains 16X16 arrays that simulate 9 channel images, hence the train_x has structure of (4608,9,16,16), the labels of this data of (4608). The validation data has dimension of (521,9,16,16) and labels of (512).

importing the libraries
import pandas as pd
import numpy as np

for reading and displying images
from skimage.io import imread
import matplotlib.pyplot as plt
%matplotlib inline

for creating validation set
from sklearn.model_selection import train_test_split

for evaluating the model
from sklearn.metrics import accuracy_score
from tqdm import tqdm

PyTorch libraries and modules
import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

check if CUDA is available
train_on_gpu = torch.cuda.is_available()
if not train_on_gpu:
print(‘CUDA is not available.  Training on CPU …’)
else:
print(‘CUDA is available!  Training on GPU …’)
train_data=images_training
test_data=images_testing
train_x = np.array(train_data)
test_x = np.array(test_data)
#agregando las etiquetas, es decir, identificando a qué clase pertenece cada imagen
train_y = np.array(label_train)
test_y = np.array(label_test)
#creando un set de validación, estos datos los tomamos de los datos de entrenamiento
train_x, val_x, train_y, val_y = train_test_split(train_x, train_y, test_size = 0.1)

convirtiendo las imagenes de ENTRENAMIENTO al formato torch para su manejo dentro de la red
#train_x tiene la forma (4608,9,16,16)
train_x  = torch.from_numpy(train_x)

convirtiendo las etiquetas de cada imagen en formato torch
train_y = train_y.astype(int);
train_y = torch.from_numpy(train_y)

estructura de los datos de entrenamiento
train_x.shape, train_y.shape

convirtiendo las imagenes de VALIDACIÓN al formato torch para su manejo dentro de la red
#train_x tiene la forma (521,9,16,16)
val_x  = torch.from_numpy(val_x)

converting the target into torch format
val_y = val_y.astype(int);
val_y = torch.from_numpy(val_y)
class Net(nn.Module):
def init(self):
super(Net, self).init()
# Defining a 2D convolution layer
    #convirtiendo la imagen de profundidad 9 a 16 con un kernel de 3x3
    self.conv1=nn.Conv2d(9, 16, kernel_size=3, stride=1, padding=1)

    #convirtiendo la imagen de profundidad 16 a 32 con un kernel de 3x3
    self.conv2=nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)

    #convirtiendo la imagen de profundidad 32 a 64 con un kernel de 3x3
    #self.conv3=nn.Conv2d(36, 72, kernel_size=3, stride=1, padding=1)

    #Capa macpooling
    self.pool=nn.MaxPool2d(2,2)

    #la dimensión de la imagen era de 16*16 después de la primera pool tengo 8*8 después de la segunda 4*4 despuès de la tercera 2*2

    self.fc1=nn.Linear(72*4*4, 256)
    self.fc2=nn.Linear(256,4)
    self.dropout=nn.Dropout(0.25)
    

# Defining the forward pass    

def forward(self, x):
    #agregando la secuencia de las capas de convolución y maxpooling
    x = self.pool(F.relu(self.conv1(x)))
    x = self.pool(F.relu(self.conv2(x)))
    #x = self.pool(F.relu(self.conv3(x)))
    #flatten, convitiéndolo a forma de vector
    x=x.view(-1,72*4*4)

    #agregando la capa de dropout para prevenir sobreajuste
    x=self.dropout(x)
    #pasando la primera capa oculta con la activación relu
    x=F.relu(self.fc1(x))
    #agregando la capa de dropout para prevenir sobreajuste
    x=self.dropout(x)
   #agregando la segunda capa oculta del fully connectec para obtener la clasificación  
    x=self.fc2(x)
    return x

#Creando la CNN
model= Net()
print (model)
#moviéndolo a una GPU en caso de que tenga CUDA disponible
if train_on_gpu:
model.cuda()
criterion = criterion.cuda()
#Definiendo que usaremos la función de perdida
criterion=nn.CrossEntropyLoss()
#especificando un optimizador con un gradiente descendiente estocástico y una
#tasa de aprendizaje de 0.01
optimizer=optim.SGD(model.parameters(),lr=0.01)
def train(epoch):
###################

# train the model #

###################

model.train()
tr_loss = 0
# getting the training set
x_train, y_train = Variable(train_x), Variable(train_y)
# getting the validation set
x_val, y_val = Variable(val_x), Variable(val_y)
# converting the data into GPU format
if torch.cuda.is_available():
    x_train = x_train.cuda()
    y_train = y_train.cuda()
    x_val = x_val.cuda()
    y_val = y_val.cuda()


# clear the gradients of all optimized variables
optimizer.zero_grad()
# forward pass: compute predicted outputs by passing inputs to the model
output_train = model(x_train.float().squeeze(1))
output_val = model(x_val.float().squeeze(1))
print('train ',x_train.shape, 'label ', y_train.shape, 'prediccion ', output_train )
print(output_train)

# computing the training and validation loss
loss_train = criterion(output_train, y_train)
print(loss_train)
loss_val = criterion(output_val, y_val)
print(loss_val)
train_losses.append(loss_train)
val_losses.append(loss_val)
#En esta sección se actualizan los pesos praa todos los modelos
# backward pass: compute gradient of the loss with respect to model parameters

loss_train.backward()
# perform a single optimization step (parameter update)
optimizer.step()
# update training loss
tr_loss = loss_train.item()
print('Epoch: {} \tTraining Loss: {:.5f} \tValidation Loss: {:.5f}'.format(
    epoch, loss_train, loss_val))


defining the number of epochs
n_epochs = 10

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)
This is the mistake
/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction)
2822     if size_average is not None or reduce is not None:
2823         reduction = _Reduction.legacy_get_string(size_average, reduce)
 → 2824     return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
2825
2826
IndexError: Target 4 is out of bounds.