start_fc error when input_dim != 1

[TonyYuanMD]

Hi, thanks for publishing the model structure on the first hand!
I wonder if it's my issue that when I use input_dim more than 1, the strat_fc(batch_x.unsqueeze(-1)) tends to mess up the dimensions in matrix multiplication?

After removing it my training ran smoothly, however it gave huge loss value. I will update after debugging and fully understanding the code :)