Flickr8k+PyTorch, CNN+LSTM predicts always same words during model testing

Question

I'm a beginner in Machine Learning and I'm working with the Flickr8k dataset (it contains ~8000 images, every image has 5 captions: ~40000 pairs). I splitted the dataset in training (70%) and validation/test (15%). I trained the model (for image captioning) on my GPU for 40 epochs (I need ~400s for each epoch).

During the model training the accuracy (n° of correct words predicted) increases but if I try to predict a caption given an image, the caption generated is always like "<start> <start> <start> <start> ..." (so my LSTM predicts always the same word/token). I think maybe the issue is the implementation of my "predict" function (since that the forward pass of the model returns a better caption each training epoch). Here's some code:

class CNNEncoder(nn.Module):
    def __init__(self, embed_size): # embed_size is the output size of the CNN encoder
        super(CNNEncoder, self).__init__()
        self.inception = models.inception_v3(weights=Inception_V3_Weights.DEFAULT)
        self.inception.fc = nn.Linear(self.inception.fc.in_features, embed_size)

        for param in self.inception.parameters():
            if param is not self.inception.fc.weight and param is not self.inception.fc.bias:
                param.requires_grad_(False)

        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
    def forward(self, input):
        if self.training:
            features, _ = self.inception(input)
        else:
            features = self.inception(input)
        return self.dropout(self.relu(features))

class DecoderRNN(nn.Module):
    def __init__(self, embed_size, hidden_size, vocab_size, n_lstm_layers=2):
        super(DecoderRNN, self).__init__()
        self.embed = nn.Embedding(vocab_size, embed_size)                    
        self.lstm = nn.LSTM(embed_size, hidden_size, num_layers=n_lstm_layers, batch_first=True)
        self.linear = nn.Linear(hidden_size, vocab_size)
        self.dropout = nn.Dropout(0.5)
    def forward(self, features, captions):
        embeddings = self.dropout(self.embed(captions))
        embeddings = torch.cat((features.unsqueeze(1), embeddings), dim=1) # shape: (batch_size, caption_length + 1, embed_size)
        outputs, _ = self.lstm(embeddings)
        return self.linear(outputs)

class ImageCaptioner(nn.Module):
    def __init__(self, embed_size, hidden_size, vocab_size, n_lstm_layers=2):
        super(ImageCaptioner, self).__init__()
        self.cnn = CNNEncoder(embed_size)
        self.decoderRNN = DecoderRNN(embed_size, hidden_size, vocab_size, n_lstm_layers)

    def forward(self, images, captions):
        features = self.cnn(images) # shape: (batch_size, embed_size)
        outputs = self.decoderRNN(features, captions) # shape: (batch_size, caption_length, vocab_size)
        return outputs
    
    def predict(self, images, device, vocab, max_length=30):
        result = torch.zeros((images.shape[0], max_length, len(vocab)), dtype=torch.float32, device=device)
        feature_vector = self.cnn(images).unsqueeze(1) 
        states = None
        for i in range(max_length):
            outputs, states = self.decoderRNN.lstm(feature_vector, states)
            outputs = self.decoderRNN.linear(outputs.squeeze(1)) 
            predicted = outputs.argmax(1) # shape: (batch_size)
            result[:, i, predicted] = 1
            
            feature_vector = self.decoderRNN.embed(predicted.unsqueeze(1)) # shape: (batch_size, 1, embed_size)
            if (predicted == vocab["<end>"]).all():
                break
            break

        return result

I already tried to change embed_size(256)/hidden_size(128)/n_lstm_layers(1) params of my model. What could be going wrong? Why does my LSTM during testing return always the same output? I already tried to test the model using the training set instead of the test set (and it always predicts only <start> tokens).