def test_rnn_summary_shapes(max_depth):
    """Test that the model summary works for RNNs."""
    model = ParityModuleRNN()

    b = 3
    t = 5
    i = model.rnn.input_size
    h = model.rnn.hidden_size
    o = model.linear_out.out_features

    model.example_input_array = torch.zeros(b, t, 10)

    summary = summarize(model, max_depth=max_depth)
    assert summary.in_sizes == [[b, t, i], [b, t, h]]  # rnn  # linear
    assert summary.out_sizes == [[[b, t, h], [[1, b, h], [1, b, h]]], [b, t, o]]  # rnn  # linear
Test that the model summary works for RNNs.
tests/tests_pytorch/utilities/test_model_summary.py