Hello
I have classification problem. My input is sequence of length 341 and output one of three classes {0,1,2}, I want to train linear regression model using Pytorch, I have the following class but during the training, the loss values start to have numbers then inf then NAN. I do not know how to fix that . Also I tried to initialize the weights for linear model but it is the same thing. Also I tried different epochs 5, 20 ,100 all are same . Any suggestions.
class regression(nn.Module):
    def __init__(self, input_dim):
        super().__init__()
        self.input_dim = input_dim
        # One layer
        self.linear = nn.Linear(input_dim, 1)

    def forward(self, x):
        y_pred = self.linear(x)
        return y_pred

criterion = torch.nn.MSELoss()

def fit(model, data_loader, optim, epochs):
    for epoch in range(epochs):

        for i, (X, y) in enumerate(data_loader):
            X = X.float()
            y = y.unsqueeze(1).float()
            X = Variable(X, requires_grad=True)
            y = Variable(y, requires_grad=True)
            # Make a prediction for the input X
            pred = model(X)
            #loss = (y-pred).pow(2).mean()
            loss = criterion(y, pred)
            optim.zero_grad()
            loss.backward()
            optim.step()
            print(loss)
            print(type(loss))
        # Give some feedback after each 5th pass through the data
        if epoch % 5 == 0:
            print("Epoch", epoch, f"loss: {loss}")
    return None
regnet = regression(input_dim=341)
optim = SGD(regnet.parameters(), lr=0.01)
fit(regnet, data_loader, optim=optim, epochs=5)
pred = regnet(torch.Tensor(test_set.data_info).float())
pred = pred.detach().numpy()