Hello, I did FNN for classification.
However, it gets stacked (kernel is dead) during the training stage.
What can be wrong with the network?
#convert to tensors
X_train_tens = torch.from_numpy(X_np)
X_train_tens = X_train_tens.type(torch.FloatTensor)
y_train_tens = torch.from_numpy(y_np)
y_train_tens = y_train_tens.type(torch.LongTensor)

device = torch.device('cpu' if torch.cuda.is_available() else 'cpu')

batch_size = 10
input_size = 154
hiden_size = 462
num_classes = 4
learning_rate = 0.001
num_epochs = 100

input_train = autograd.Variable(X_train_tens)
target_train = autograd.Variable(y_train_tens)

#define FNN

class Net(nn.Module):
    def __init__(self,input_size,hiden_size,num_classes):
        super().__init__()
        self.fc1 = nn.Linear(input_size, hiden_size)
        self.fc2 =nn.Linear(hiden_size, num_classes)
    
    def forward(self, x):
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        return x

model = Net(input_size=input_size, hiden_size=hiden_size, num_classes=num_classes).to(device)
opt = torch.optim.Adam(params=model.parameters(), lr=learning_rate)
loss_func = torch.nn.CrossEntropyLoss()

#run training stage

correct = 0
total = 0
counter = 0
for epoch in range (num_epochs):
    out = model(input_train).to(device)
    _, pred = out.max(1)
    total += target_train.size(0)
    correct += (pred == target_train).sum().item()
    print(input_train)
    print(pred)
    loss = loss_func(out,target_train)
    counter +=1
    print('loss train', "Epoch N", counter,loss.data[0])
    model.zero_grad()
    loss.backward()
    opt.step()
print('Accuracy of the network on train dataset: {} %'.format(100 * correct / total))