Description

Add an optional batch_size_fn: Callable[[list], int] to the batch decorator to allow different ways to determine "batch size".

Use case

In some scenarios (like Graph Neural nets for me) batching is more effective by setting a maximum number of total nodes instead of using total number of examples/graphs. Allowing the user to pass a optional function to determine an effective batch size would be great!

For example I would like use something like batch_size_fn = lambda batch: sum(graph.num_nodes for graph in batch); and have that be used instead of len(batch) here.