Your code is not executable, but I guess U
has the wrong shape:
D_m = 2024
D_e = 1
dropout=0.5
lstm = nn.LSTM(input_size=D_m, hidden_size=D_e, num_layers=2, bidirectional=True, dropout=dropout)
x = torch.randn(10, 10, D_m)
out = lstm(x)
print(out[0].shape)
# torch.Size([10, 10, 2])
# error
x = torch.randn(10, 10, 100)
out = lstm(x)
# RuntimeError: input.size(-1) must be equal to input_size. Expected 2024, got 100