def test_dynamic_batch_sizes(self, block_constructor):
    """Checks that all batch sizes are as expected through a GraphNetwork."""
    input_graph = self._get_input_graph()
    placeholders = input_graph.map(_mask_leading_dimension, graphs.ALL_FIELDS)
    model = block_constructor(
        functools.partial(snt.nets.MLP, output_sizes=[10]))
    output = model(placeholders)
    with tf.Session() as sess:
      sess.run(tf.global_variables_initializer())
      other_input_graph = utils_np.data_dicts_to_graphs_tuple(
          [SMALL_GRAPH_1, SMALL_GRAPH_2])
      actual = sess.run(output, {placeholders: other_input_graph})
    for k, v in other_input_graph._asdict().items():
      self.assertEqual(v.shape[0], getattr(actual, k).shape[0])