Predictions being a straight line along the x-axis

[toibanoor]

I have implemented the example shared, but the doesn't seem to learn and the loss isn't decreasing.
I used the following training loop.
model = tst.TimeSeriesTransformer(
input_size=1,
dec_seq_len=enc_seq_len,
batch_first=batch_first,
num_predicted_features=1
)

Define your loss function and optimizer

criterion = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

Number of training epochs

num_epochs = 10

for epoch in range(num_epochs):
model.train()
for i, batch in enumerate(training_data):

    src, trg, trg_y = batch
    

    # Zero the gradients
    optimizer.zero_grad()

    # Permute from shape [batch size, seq len, num features] to [seq len, batch size, num features]
    if batch_first == False:

        shape_before = src.shape
        src = src.permute(1, 0, 2)
        #print("src shape changed from {} to {}".format(shape_before, src.shape))

        shape_before = trg.shape
        trg = trg.permute(1, 0, 2)
        #print("trg shape changed from {} to {}".format(shape_before, trg.shape))
    

    output = model(
        src=src,
        tgt=trg,
        src_mask=src_mask,
        tgt_mask=tgt_mask)
        
 
    
    # # Permute output from [seq len, batch size] to [batch size, seq len]
    trg_y = trg_y.reshape(trg_y.shape[0], trg_y.shape[1], 1)
    trg_y = trg_y.permute(1, 0, 2)

    # print('output shape: ', output.shape)
    # print('target_y shape: ', trg_y.shape)



    # Calculate the loss
    loss = criterion(output, trg_y)

    # Backpropagation
    loss.backward()

    # Update parameters
    optimizer.step()

# Print training progress
print(f'Epoch [{epoch+1}/{num_epochs}], Batch [{i+1}/{len(training_data)}], Loss: {loss.item()}')

Could you please tell me what I am doing wrong?