BLEURT on GPU

#3
by malper - opened

Is it possible to calculate BLEURT on a specified device such as a GPU using this API?
If so, it would be helpful to add this information to the README.
It is also not clear how to adjust the batch size of inputs used for this calculation, which is important when calculating BLEURT scores on a large dataset.

Evaluate Metric org

Hi @malper

BLEURT uses TensorFlow under the hood which should take care of device placement (it should prioritize GPUs if available). As for BS this is not currently implemented. The score (see here) method would take an additional kwarg for the batch_size which defaults to 16.

Sign up or log in to comment