# How to Build LSTM or GRU for Sequence-to-Sequence Tasks in PyTorch vs TensorFlow

> Compare PyTorch and TensorFlow for building LSTM or GRU sequence-to-sequence models. Learn control over hidden states and faster prototyping with these deep learning frameworks.

- Repository: [Microsoft/AI-For-Beginners](https://github.com/microsoft/AI-For-Beginners)
- Tags: tutorial
- Published: 2026-08-20

---

**Use PyTorch for fine-grained control over hidden states and custom training loops, or TensorFlow/Keras for declarative state propagation and faster prototyping with `return_state=True` and `initial_state` parameters.**

Sequence-to-sequence (seq2seq) models map an input token sequence to an output token sequence using an **encoder-decoder architecture**. Both the encoder and decoder can be implemented with **LSTM** or **GRU** recurrent cells. This guide walks through building these models in PyTorch and TensorFlow using the reference implementations from Microsoft's [AI-For-Beginners](https://github.com/microsoft/AI-For-Beginners) repository.

## Core Architecture: Encoder-Decoder Pattern

A seq2seq model consists of two components:

- **Encoder**: Processes the source sequence and compresses it into a hidden state vector
- **Decoder**: Unfolds the hidden state to generate the target sequence token-by-token

The encoder's final hidden state becomes the decoder's initial hidden state. During training, **teacher forcing** feeds ground-truth tokens to the decoder; during inference, the decoder uses its own predictions.

## PyTorch Implementation: Manual State Management

In PyTorch, you instantiate `torch.nn.LSTM` or `torch.nn.GRU` directly and manage hidden states explicitly.

### Cell Definition and Return Values

```python
import torch
import torch.nn as nn

# LSTM returns (output, (h_n, c_n)) — hidden state AND cell state

lstm = nn.LSTM(input_size=256, hidden_size=512, batch_first=True)

# GRU returns (output, h_n) — hidden state only

gru = nn.GRU(input_size=256, hidden_size=512, batch_first=True)

```

### Building the Encoder

The encoder in [`lessons/5-NLP/16-RNN/torchnlp.py`](https://github.com/microsoft/AI-For-Beginners/blob/main/lessons/5-NLP/16-RNN/torchnlp.py) demonstrates the LSTM pattern. For seq2seq, extract the final hidden state:

```python
class Encoder(nn.Module):
    def __init__(self, vocab_sz, emb_dim, hid_dim, cell='LSTM'):
        super().__init__()
        self.embedding = nn.Embedding(vocab_sz, emb_dim)
        rnn_cls = nn.LSTM if cell == 'LSTM' else nn.GRU
        self.rnn = rnn_cls(emb_dim, hid_dim, batch_first=True)

    def forward(self, src):
        emb = self.embedding(src)                # (batch, src_len, emb_dim)

        outputs, hidden = self.rnn(emb)          # hidden = (h_n, c_n) for LSTM

        return hidden

```

### Building the Decoder with Teacher Forcing

```python
class Decoder(nn.Module):
    def __init__(self, vocab_sz, emb_dim, hid_dim, cell='LSTM'):
        super().__init__()
        self.embedding = nn.Embedding(vocab_sz, emb_dim)
        rnn_cls = nn.LSTM if cell == 'LSTM' else nn.GRU
        self.rnn = rnn_cls(emb_dim, hid_dim, batch_first=True)
        self.fc = nn.Linear(hid_dim, vocab_sz)

    def forward(self, trg, hidden):
        emb = self.embedding(trg)                # (batch, 1, emb_dim)

        output, hidden = self.rnn(emb, hidden)   # one step at a time

        pred = self.fc(output.squeeze(1))        # (batch, vocab_sz)

        return pred, hidden

```

### Handling Variable-Length Sequences

PyTorch provides `torch.nn.utils.rnn.pack_padded_sequence` and `pad_packed_sequence` for efficient batching:

```python
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

# Sort by length descending

lengths = [len(seq) for seq in sequences]
packed = pack_padded_sequence(padded, lengths, batch_first=True)
output, hidden = self.rnn(packed)
output, _ = pad_packed_sequence(output, batch_first=True)

```

### Complete Training Loop

```python
encoder = Encoder(VOCAB_SIZE, EMB_DIM, HID_DIM, cell='GRU')
decoder = Decoder(VOCAB_SIZE, EMB_DIM, HID_DIM, cell='LSTM')
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(list(encoder.parameters()) + list(decoder.parameters()))

for epoch in range(10):
    optimizer.zero_grad()
    src = torch.randint(0, VOCAB_SIZE, (32, 15))
    trg = torch.randint(0, VOCAB_SIZE, (32, 10))

    enc_hidden = encoder(src)
    dec_input = trg[:, 0].unsqueeze(1)
    loss = 0
    hidden = enc_hidden
    
    for t in range(1, trg.size(1)):
        pred, hidden = decoder(dec_input, hidden)
        loss += criterion(pred, trg[:, t])
        dec_input = trg[:, t].unsqueeze(1)  # teacher forcing

    
    loss.backward()
    optimizer.step()

```

The [`torchnlp.py`](https://github.com/microsoft/AI-For-Beginners/blob/main/torchnlp.py) file shows how to swap LSTM for GRU with a single line change: replace `nn.LSTM` with `nn.GRU`.

## TensorFlow/Keras Implementation: Declarative State Flow

TensorFlow/Keras handles state propagation automatically through `return_state=True` and `initial_state` parameters.

### Cell Definition with State Returns

```python
import tensorflow as tf
from tensorflow.keras import layers

# With return_state=True: returns (output, h, c) for LSTM

lstm = layers.LSTM(units=512, return_state=True, name='encoder')

# With return_state=True: returns (output, h) for GRU

gru = layers.GRU(units=512, return_state=True, name='encoder')

```

### Building the Encoder

The `RNNTF.ipynb` notebook demonstrates the `return_state=True` pattern essential for seq2seq encoders:

```python
class Encoder(tf.keras.layers.Layer):
    def __init__(self, vocab_sz, emb_dim, hid_dim, cell='LSTM'):
        super().__init__()
        self.embedding = layers.Embedding(vocab_sz, emb_dim)
        rnn_cls = layers.LSTM if cell == 'LSTM' else layers.GRU
        self.rnn = rnn_cls(hid_dim, return_state=True, name='encoder')

    def call(self, x):
        x = self.embedding(x)
        if isinstance(self.rnn, layers.LSTM):
            output, state_h, state_c = self.rnn(x)
            return [state_h, state_c]  # list for LSTM

        else:
            output, state_h = self.rnn(x)
            return [state_h]            # list for GRU

```

### Building the Decoder with Initial State

```python
class Decoder(tf.keras.layers.Layer):
    def __init__(self, vocab_sz, emb_dim, hid_dim, cell='LSTM'):
        super().__init__()
        self.embedding = layers.Embedding(vocab_sz, emb_dim)
        rnn_cls = layers.LSTM if cell == 'LSTM' else layers.GRU
        self.rnn = rnn_cls(hid_dim, return_sequences=True, 
                          return_state=True, name='decoder')
        self.fc = layers.Dense(vocab_sz)

    def call(self, x, states):
        x = self.embedding(x)
        if isinstance(self.rnn, layers.LSTM):
            output, state_h, state_c = self.rnn(x, initial_state=states)
            states = [state_h, state_c]
        else:
            output, state_h = self.rnn(x, initial_state=states)
            states = [state_h]
        logits = self.fc(output)
        return logits, states

```

### Complete Training with GradientTape

```python
encoder = Encoder(VOCAB_SIZE, EMB_DIM, HID_DIM, cell='GRU')
decoder = Decoder(VOCAB_SIZE, EMB_DIM, HID_DIM, cell='LSTM')

loss_obj = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
optimizer = tf.keras.optimizers.Adam()

with tf.GradientTape() as tape:
    enc_states = encoder(src)
    dec_input = tf.expand_dims(trg[:, 0], 1)
    loss = 0
    states = enc_states
    
    for t in range(1, trg.shape[1]):
        logits, states = decoder(dec_input, states)
        loss += loss_obj(trg[:, t], tf.squeeze(logits, axis=1))
        dec_input = tf.expand_dims(trg[:, t], 1)  # teacher forcing

grads = tape.gradient(loss, encoder.trainable_variables + decoder.trainable_variables)
optimizer.apply_gradients(zip(grads, encoder.trainable_variables + decoder.trainable_variables))

```

## Key Differences: PyTorch vs TensorFlow for Seq2Seq

| Aspect | PyTorch | TensorFlow/Keras |
|--------|---------|------------------|
| **State handling** | Manual tuple unpack `(h_n, c_n)` vs `h_n` | Automatic via `return_state` and `initial_state` |
| **Cell switching** | Change `nn.LSTM` to `nn.GRU` — same API | Change `layers.LSTM` to `layers.GRU` — handles state count automatically |
| **Variable lengths** | `pack_padded_sequence` / `pad_packed_sequence` | Padded batches with `layers.Masking` |
| **Training loop** | Explicit `for` loop with manual loss accumulation | `model.fit()` or `GradientTape` with similar loop structure |
| **Flexibility** | Easier custom attention, beam search | Encapsulated higher-level APIs |

## Advanced Patterns from the Repository

### Bidirectional Encoders

Both frameworks support bidirectional encoders for richer context:

**PyTorch**:

```python
self.rnn = nn.LSTM(emb_dim, hid_dim, batch_first=True, bidirectional=True)

# Output is 2*hid_dim, handle concatenation in decoder

```

**TensorFlow/Keras**:

```python
self.rnn = layers.Bidirectional(layers.LSTM(hid_dim, return_state=True))

# Returns forward and backward states; concatenate or sum as needed

```

The `RNNTF.ipynb` file shows bidirectional LSTM implementation for text classification.

### Multi-Layer Stacking

The `NER-TF.ipynb` notebook implements a two-layer bidirectional LSTM encoder:

```python
layers.Bidirectional(layers.LSTM(HID_DIM, return_sequences=True, return_state=True))

```

For PyTorch, set `num_layers=2`:

```python
nn.LSTM(emb_dim, hid_dim, num_layers=2, batch_first=True, bidirectional=True)

```

### TimeDistributed Decoders

The NER notebook uses `TimeDistributed` dense layers for per-token predictions:

```python
layers.TimeDistributed(layers.Dense(num_tags, activation='softmax'))

```

In PyTorch, this is implicit — the linear layer operates on all timesteps when given the full sequence.

## Inference: Greedy Decoding Without Teacher Forcing

**PyTorch**:

```python
encoder.eval()
decoder.eval()
with torch.no_grad():
    hidden = encoder(src)
    dec_input = torch.full((batch_size, 1), START_ID)
    for _ in range(max_len):
        pred, hidden = decoder(dec_input, hidden)
        dec_input = pred.argmax(dim=-1).unsqueeze(1)

```

**TensorFlow/Keras**:

```python
def greedy_decode(encoder, decoder, src_seq, max_len=20, start_id=1):
    enc_states = encoder(src_seq)
    dec_input = tf.expand_dims([start_id] * src_seq.shape[0], 1)
    outputs = []
    states = enc_states
    for _ in range(max_len):
        logits, states = decoder(dec_input, states)
        pred_id = tf.argmax(logits, axis=-1, output_type=tf.int32)
        outputs.append(pred_id)
        dec_input = pred_id
    return tf.concat(outputs, axis=1)

```

## Reference Files in AI-For-Beginners

| File | Description |
|------|-------------|
| [`lessons/5-NLP/16-RNN/torchnlp.py`](https://github.com/microsoft/AI-For-Beginners/blob/main/lessons/5-NLP/16-RNN/torchnlp.py) | PyTorch LSTM/GRU classifier with pack-sequence utilities |
| `lessons/5-NLP/16-RNN/RNNTF.ipynb` | TensorFlow bidirectional LSTM with `return_state=True` pattern |
| `lessons/5-NLP/19-NER/NER-TF.ipynb` | Full seq2seq with two-layer bidirectional LSTM encoder and `TimeDistributed` decoder |
| [`environment.yml`](https://github.com/microsoft/AI-For-Beginners/blob/main/environment.yml) | Dependency versions for TensorFlow, PyTorch, and CUDA |

## Summary

- **PyTorch** requires explicit handling of hidden-state tuples `(h_n, c_n)` for LSTM or `h_n` for GRU, offering fine-grained control for custom attention and beam search
- **TensorFlow/Keras** automates state propagation through `return_state=True` and `initial_state`, reducing boilerplate for standard architectures
- Both frameworks support **bidirectional** encoders, **multi-layer** stacking, and **mixed cell types** (LSTM encoder + GRU decoder or vice versa)
- The [`torchnlp.py`](https://github.com/microsoft/AI-For-Beginners/blob/main/torchnlp.py) file demonstrates single-line cell swapping; `NER-TF.ipynb` shows production-ready seq2seq patterns

## Frequently Asked Questions

### Can I mix LSTM and GRU between encoder and decoder?

Yes. In PyTorch, initialize `encoder.rnn = nn.LSTM(...)` and `decoder.rnn = nn.GRU(...)`, then handle the state mismatch — GRU's single state `h_n` becomes both `h_0` and `c_0` for LSTM (or initialize `c_0` to zeros). In TensorFlow, the `initial_state` parameter accepts lists of appropriate length automatically.

### How do I handle variable-length sequences in batches?

Use PyTorch's `pack_padded_sequence` and `pad_packed_sequence` for computation efficiency. In TensorFlow, provide padded tensors and use `layers.Masking` or the `mask` parameter; Keras automatically propagates masks through compatible layers.

### Which is faster for seq2seq training: PyTorch or TensorFlow?

Performance is comparable when using equivalent implementations. PyTorch's dynamic graphs allow easier debugging and modification. TensorFlow's `tf.function` compilation and XLA can optimize static graphs for production deployment. For research flexibility, prefer PyTorch; for production pipelines, consider TensorFlow's SavedModel format.

### Where can I find working attention mechanisms for these models?

The AI-For-Beginners repository focuses on foundational RNNs. For attention, extend the decoder to compute context vectors using `torch.bmm` (PyTorch) or `layers.Attention` (TensorFlow). The `torch.nn.MultiheadAttention` and `layers.MultiHeadAttention` modules provide optimized implementations for transformer-style attention over recurrent outputs.