import torch
import torch.nn as nn

# Make a GRU: input size 1, hidden size 1, single layer
gru = nn.GRU(input_size=1, hidden_size=1, batch_first=True)

# Sequence: batch size 1, sequence length 5, 1 feature per step
input_seq = torch.tensor([[[1.0], [2.0], [3.0], [4.0], [5.0]]])  # shape (1, 5, 1)

# Initial hidden state: zeros
h0 = torch.zeros(1, 1, 1)

# Forward pass
output_seq, hn = gru(input_seq, h0)

print("Output at each timestep:")
print(output_seq)
print("Final hidden state:")
print(hn)