Neural Network for Regression using PyTorch
I am trying to implement a Neural Network for predicting the h1_hemoglobin in PyTorch. After creating a model, I kept 1 in the output layer as this is Regression. But I got the error as below. I'm not able to understand the mistake. Keeping a large value like 100 in the output layer removes the error but renders the model useless as I am trying to implement regression.
Data:

from sklearn.model_selection import train_test_split
X_train,X_test,y_train,y_test=train_test_split(X,y,test_size=0.2,random_state=0)
##### Creating Tensors
X_train=torch.tensor(X_train)
X_test=torch.tensor(X_test)
y_train=torch.LongTensor(y_train)
y_test=torch.LongTensor(y_test)

class ANN_Model(nn.Module):
    def __init__(self,input_features=4,hidden1=20,hidden2=20,out_features=1):
        super().__init__()
        self.f_connected1=nn.Linear(input_features,hidden1)
        self.f_connected2=nn.Linear(hidden1,hidden2)
        self.out=nn.Linear(hidden2,out_features)
    def forward(self,x):
        x=F.relu(self.f_connected1(x))
        x=F.relu(self.f_connected2(x))
        x=self.out(x)
        return x

loss_function = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr = 0.01)

epochs = 500
final_losses = []
for i in range(epochs):
    i = i + 1
    y_pred = model.forward(X_train.float())
    loss=loss_function(y_pred, y_train)
    final_losses.append(loss.item())
    if i%10==1:
        print(&quot;Epoch number: {} and the loss: {}&quot;.format(i, loss.item()))
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Error:


Since you are performing regression, the CrossEntropyLoss() internally implements the NLLLoss() function. The CrossEntropyLoss() expects C classes for C predictions but you have specified only one class. The NLLLoss() tries to index into the prediction logits based on the ground-truth value. E.g., in your case, the ground-truth is a single value 14. The loss step tries to index into the 14th logit of your predictions to get its corresponding value so that it can compute the negative log likelihood on it, which is essentially - -log(probability_k) where k is the index that the ground-truth outputs. Since you have only logit in your predictions, it throws an error - index out of bounds.
For regression problems, you should consider using distance based losses such as MSELoss().
Try replacing your loss function - loss_function = CrossEntropyLoss() with loss_function = MSELoss()
