So I was able to solve the embedding dimension issue. For the training, I have written this:
for epoch in range(epochs):
    loss = 0
    train_model.train()
'''Returns the signal to construct and embeddings'''
    for encoded_data in tqdm(train_loader): 
        signal = signal.squeeze()

'''Since the embeddings are of different shape (128,10) or (499,10), I parse a batch of 10 embeddings'''

        for i in range(0,encoded_data.shape[1],10):

            signal_train = torch.reshape()
                
            input_data = encoded_data[:,i:i+10,:].squeeze().cuda()
            signal_train = signal_train.cuda()
            optimizer.zero_grad()
            signal_recon = train_model(input_data)

            loss = criterion(signal_recon, signal_train)
            loss.backward()
            
            optimizer.step()
            
            loss += loss.item()

So I wanted to know if I am doing this correctly or you recommend any changes for faster training or more accurate training?