TypeError: Got unsupported ScalarType BFloat16 with OCRErrorPredictor

[kevinhu]

I'm running Surya via Marker on a GPU, and I'm attempting to use the bfloat16 datatype. However, I'm running into this TypeError because of the order in which the cast is performed:

surya/surya/ocr_error/__init__.py

Lines 47 to 50 in aa8ee5a

	with torch.inference_mode():
	pred = self.model(batch_input_ids, attention_mask=batch_attention_mask)
	logits = pred.logits.detach().cpu().numpy().astype(np.float32)
	predictions.extend(np.argmax(logits, axis=1).tolist())

This can be fixed by performing the cast before the tensors are moved to NumPy:

logits = pred.logits.to(torch.float32).detach().cpu().numpy()

When running on GPU, this should also be faster than casting in NumPy.

[VikParuchuri]

We usually use float16 for this model - any specific reason to use bfloat?

[kevinhu]

In a couple of our benchmarks, we saw that bfloat16 gives better accuracy than float16—it seems that there's some precision errors with float16.