LSTMLastStepBlock

class maze.perception.blocks.joint_blocks.lstm_last_step.LSTMLastStepBlock(*args: Any, **kwargs: Any)

A block containing a LSTM perception block followed by a Slicing Block keeping only the output of the final time step.

For details on flattening see LSTMBlock. For details on dense layers see SliceBlock.

Parameters
  • in_keys – One key identifying the input tensors.

  • out_keys – One key identifying the output tensors.

  • in_shapes – List of input shapes.

  • hidden_size – The number of features in the hidden state.

  • num_layers – Number of recurrent layers.

  • bidirectional – If True, becomes a bidirectional LSTM.

  • non_lin – The non-linearity to apply after the final layer.

forward(block_input: Dict[str, torch.Tensor]) → Dict[str, torch.Tensor]

(overrides PerceptionBlock)

implementation of PerceptionBlock interface