# Question on training LSTM to play Snake

**URL:** <https://discuss.pytorch.org/t/question-on-training-lstm-to-play-snake/80006>\
**Category:** reinforcement-learning\
**Created:** [May 6, 2020, 1:21pm UTC](https://discuss.pytorch.org/t/question-on-training-lstm-to-play-snake/80006 "2020-05-06T13:21:10Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![waw\_waw](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/waw_waw/32/16232_2.png) [@waw\_waw](https://discuss.pytorch.org/u/waw_waw)\
**Post date:** [May 6, 2020, 1:21pm UTC](https://discuss.pytorch.org/t/question-on-training-lstm-to-play-snake/80006/1 "2020-05-06T13:21:10Z")

</div>

Hi,

I follow the DQN tutorial and trained a CNN to play a game, now I want to switch to LSTM but got a problem.

I found the tutorial for LSTM, and it recommends using (time sequence, Batch size, features). In my code, my input to LSTM is (4,32,100), where 4 means 4 consecutive frames, 32 is the batch size, 100 is a vector representing the current state.

Then I add a nn.Linear layer after LSTM, and the input size of linear layer is 4_32_lstm\_output\_size, and here comes the problem. In trainig, the batch size is 32-\>(4,32,100), but in testing, the batch size is1 -\> (4,1,100), which will cause an error.

I tried to train the LSTM with batch size 1, but it will take significantly longer time, is there any way that can let me train the LSTM with batch=32 and do inference with batch=1?

---

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [May 7, 2020, 5:10am UTC](https://discuss.pytorch.org/t/question-on-training-lstm-to-play-snake/80006/2 "2020-05-07T05:10:31Z")

</div>

The number seems to be formatted a bit wrong, but I understand that the input features to the linear layer are defined as `seq_len * batch_size * lstm_output_size`?  
You shouldn’t use the batch size as the number of input features, as this will limit your use case to only this particular size, and assumes that the samples in the batch are somehow related to each other.

Instead, `permute` the LSTM output so that the batch size is at dim0 and reset the `in_features` to e.g. `seq_len * lstm_output_size`.  
This will allow you to use arbitrary batch sizes.

Alternatively, you could also use `batch_first=True` when you are creating an instance of `LSTM`, which will accept and return the batch size in dim0.