Fix reshape LSTM states

Description

fixes #111 and fixes some warnings in the tests

@kolbytn could you review that PR?

[x] Bug fix (non-breaking change which fixes an issue)
[ ] New feature (non-breaking change which adds functionality)
[ ] Breaking change (fix or feature that would cause existing functionality to change)
[ ] Documentation (update in the documentation)

[x] I've read the CONTRIBUTION guide (required)
[ ] The functionality/performance matches that of the source (required for new training algorithms or training-related features).
[ ] I have updated the tests accordingly (required for a bug fix or a new feature).
[ ] I have included an example of using the feature (required for new features).
[ ] I have included baseline results (required for new training algorithms or training-related features).
[ ] I have updated the documentation accordingly.
[x] I have updated the changelog accordingly (required).
[x] I have reformatted the code using make format (required)
[ ] I have checked the codestyle using make check-codestyle and make lint (required)
[ ] I have ensured make pytest and make type both pass. (required)

Note: we are using a maximum length of 127 characters per line