Dropout is not working in train mode on

Hello,
I am trying to do Monte Carlo sampling to estimate uncertainty. Even though I turn on the train mode by model.train(), dropout is not working and all the samples are getting the same value. Can anybody help me to identify what am I doing wrong here? Thanks in advance!

Model looks like something like below:

LSTM_Base(
  (lstm): LSTM(167, 256, num_layers=3, batch_first=True, dropout=0.5)
  (fc): Linear(in_features=12800, out_features=256, bias=True)
  (dropout): Dropout(p=0.5, inplace=False)
  (fc2): Linear(in_features=256, out_features=43, bias=True)
)

-My sampling function looks like something like this

def sampler(X, model, T=128):
    y_dist = []
    model.train()
    hidden_init = (torch.zeros(n_layers, 1, hidden_dim).cuda(), 
                       torch.zeros(n_layers, 1, hidden_dim).cuda())
    if model.training:
        for i in range(T):
            sample_out, _ = model.train().forward(X, hidden_init)
            y_dist.append(torch.clone(sample_out))
    
    return y_dist
  • Note that model is already trained, I am trying to do sampling in train mode, after I trained the model.

EDIT : I just realized that there is no problem about the working of train mode. Because before training, the model is able to predict different outputs in train mode. But after training with more than 10 epochs, all outputs become the same. After realizing that, I also tried to train for just one epoch, and see that outputs are different. So it is clear that the problem is not about train mode. Can it be vanishing gradients? I am looking forward to read you suggestions!