pyprobound.utils.get_split_size

get_split_size(max_embedding_size, max_split, device)

Calculates the minibatch needed to avoid GPU limits on tensor sizes.

Parameters:

• max_embedding_size (int) – The maximum number of bytes needed to encode a sequence.

• max_split (int) – Maximum number of sequences scored at a time (lower values reduce memory but increase computation time).

• device (str | device) – The current device of the model.

Return type:

int