I have 500 states, actions, rewards, and hidden states. But I get the runtime error because of mismatched shapes.

Exception has occurred: RuntimeError The size of tensor a (500) must match the size of tensor b (16) at non-singleton dimension 0 File "/mnt/c/Users/User/code/VideoPPO/test.py", line 33, in outputs = model(features, (h_x, c_x)) RuntimeError: The size of tensor a (500) must match the size of tensor b

(16) at non-singleton dimension 0

It happens because I have 500 states, but after splitting them into sequences with 32 states each, we need 32 times fewer hidden states and I don't know how to do that. That's we have hidden states with the shape [2, 500, 256], but we need shape [2, 16, 256].

A small clarification: this is simplified code for an RL task where we concatenate states, actions, and rewards, then send them into an LSTM for results.

There's code.

import torch as t
from torch import nn

model = nn.LSTM(67, 256, num_layers=2, batch_first=True)

states = t.randn(500, 64) # [n_steps, state_dim]
actions = t.randn(500, 2) # [n_steps, action_dim]
rewards = t.randn(500, 1) # [n_steps, reward_dim]

h_x = t.randn(2, 500, 256) # [num_layers, n_steps, hidden_size]
c_x = t.randn(2, 500, 256) # [num_layers, n_steps, hidden_size]

# Split into sequences.
# We get 16 sequences where the last one is not full.
states = t.split(states, split_size_or_sections=32)
actions = t.split(actions, split_size_or_sections=32)
rewards = t.split(rewards, split_size_or_sections=32)

# Compute unpadded lengths
lengths = [i.shape[0] for i in states]

# Padding
states = t.nn.utils.rnn.pad_sequence(states, batch_first=True, padding_value=0)
actions = t.nn.utils.rnn.pad_sequence(actions, batch_first=True, padding_value=0)
rewards = t.nn.utils.rnn.pad_sequence(rewards, batch_first=True, padding_value=0)

# Concatenation and packing.
features = t.cat([states, actions, rewards], dim=-1)
features = t.nn.utils.rnn.pack_padded_sequence(features, lengths, batch_first=True, enforce_sorted=True)

# LSTM forward
with t.no_grad():
    """
    Exception has occurred: RuntimeError
    The size of tensor a (500) must match the size of tensor b (16) at non-singleton dimension 0
    File "/mnt/c/Users/User/code/VideoPPO/test.py", line 33, in 
        outputs = model(features, (h_x, c_x))
    RuntimeError: The size of tensor a (500) must match the size of tensor b (16) at non-singleton dimension 0
    """
    """
    Error has occured because we have hidden states with the shape
    [2, 500, 256] while as we must have shape [2, 500 // 32, 256] and we need 32 times fewer of data.
    """
    outputs = model(features, (h_x, c_x)) 

# Unpacking
outputs, lengths = t.nn.utils.rnn.pad_packed_sequence(features, batch_first=True, padding_value=0.0)

# Unpadding
outputs = t.nn.utils.rnn.unpad_sequence(outputs, lengths, batch_first=True)

# Concatenation
outputs = t.cat(outputs, dim=0)

pass

Or should I do it in a completely different way?