I am building a seq2seq model which should be able to mark the start and end indexes (e.g. “step start” and “step end”) on the sensor data attached to an athlete’s leg. I would like to do the same for kicks and turns.

However, I am having difficulty producing a (100, 3) matrix using an LSTM model. Is an approach similar to the following code example appropriate for this problem?

```
class StepModel(nn.Module):
def __init__(self, input_size, hidden_size, num_layers, output_size):
super(StepModel, self).__init__()
self.hidden_size = hidden_size
self.num_layers = num_layers
self.lstm = nn.LSTM(input_size, hidden_size, num_layers, bidirectional=True, batch_first=True)
self.fc = nn.Linear(hidden_size, output_size)
self.softmax = nn.Softmax(dim=-1)
def forward(self, x):
h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
out, _ = self.lstm(x, (h0, c0))
out = self.fc(out[:, -1, :])
out = self.softmax(out)
return out
```