The input size inside my LSTM model inside the forward pass looks like this :
input :torch.Size([1024, 1, 1])

I think I have a problem dealing with batch, for more information I included my training function:
def train(model, device, federated_train_loader, optimizer, epoch):
    model.train()
    # Iterate through each gateway's dataset
    for idx, (seq, labels) in enumerate(federated_train_loader):
        batch_idx = idx+1
        # Send the model to the right gateway
        model.send(seq.location)
        # Move the data and target labels to the device (cpu/gpu) for computation
        seq, labels = seq.to(device), labels.to(device)
        # Clear previous gradients (if they exist)
        optimizer.zero_grad()
        # Make a prediction
        print('seq shape : {}'.format(seq.shape))
        print('labels shape : {}'.format(labels.shape))
        output = model(seq)
        # Calculate huber loss for regression problems
        #labels =labels.view(-1)
        #seq = seq.view(-1)
        #labels = labels.unsqueeze(1)
        #labels = labels.float()
        loss = loss_function(output, labels)
        # Calculate the gradients
        loss.backward()
        # Update the model weights
        optimizer.step()
        # Get the model back from the gateway
        #model.get()
        if batch_idx==len(federated_train_loader) or (batch_idx!=0 and batch_idx % LOG_INTERVAL == 0):
            # get the loss back
            loss = loss.get()
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * BATCH_SIZE, len(federated_train_loader) * BATCH_SIZE,
                100. * batch_idx / len(federated_train_loader), loss.item()))