After some iterations pytorch outputs same outputs

I am currently trying to implement a neural network with pytorch but while doing the training after some iterations my neural network starts giving same outputs. My nn has a size of 15 features and outputs a number.

*As a small edit after running the tests I noticed that this problem occurs whenever I run optimiser.step() in for loop.

tensor([[1.5137],
        [1.5137],
        [1.5137],
        [1.5137],
        [1.5137],
        [1.5137],
        [1.5137],
        [1.5137],
        [1.5137],
        [1.5137]], grad_fn=<AddmmBackward0>)

This is my output from my third iteration. what might be causing the issue?

class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.fc1 = nn.Linear(15, 10)
        self.fc2 = nn.Linear(10, 5)
        self.fc3 = nn.Linear(5, 1)
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
net = NeuralNetwork()

import torch.optim as optim
criterion = nn.MSELoss()
optimizer = optim.SGD(net.parameters(), lr=0.01)

batch_size=10

for i in range(0, train_X.size()[0], batch_size):
    X_batch=train_X[i:i+10].view(-1, 15)
    #print(X_batch.size())
    y_batch = train_y[i:i+10].view(-1, 1)
    print(y_batch.size())
    optimizer.zero_grad()
    output = net(train_X)
    print(f"output size is:{output.size()}")
    print(output)
    #print(output)
    loss = criterion(output, y_batch)
    loss.backward()
    optimizer.step()

from Recent Questions - Stack Overflow https://ift.tt/3r1cWty
https://ift.tt/eA8V8J