Linear layer on a 2D time series data

1

For a particular application, I am porting the code from keras to pytorch. The input is of the size [ bs x timesteps x features ], lstm output is [ bs x time step x hidden ]. Now I want to reduce this to [ bs x time step x out_features](time distributed layer on keras)
Using linear,

nn.Linear(in_features=hidden, out_features=out_features)

Is this the right way to do this if I want to preserve time information or do I need to reshape the data using contiguous in any way to achieve it?

Any help appreciated.

2

For anyone who might need help on this in the future, the linear layer as mentioned in the post works correctly. Using MSE loss on such data, works normally while using a cross entrophy loss, there might be an issue.
On a time step x feats data, I did not find a way to state the dim to apply softmax on. Instead you can specify it in log softmax and then use nllloss. (you get the same loss overall with proper softmax norm)