def test_linear_model_summary_shapes(device_str, max_depth):
    """Test that the model summary correctly computes the input- and output shapes."""
    device = torch.device(device_str)
    model = UnorderedModel().to(device)
    model.train()
    summary = summarize(model, max_depth=max_depth)
    assert summary.in_sizes == [[2, 10], [2, 7], [2, 3], [2, 7], UNKNOWN_SIZE]  # layer 2  # combine  # layer 1  # relu
    assert summary.out_sizes == [[2, 2], [2, 9], [2, 5], [2, 7], UNKNOWN_SIZE]  # layer 2  # combine  # layer 1  # relu
    assert model.training
    assert model.device == device
Test that the model summary correctly computes the input- and output shapes.
tests/tests_pytorch/utilities/test_model_summary.py