Inferring super large protein with reduced accuracy

[Ian-pku]

Hi!
I'm trying to infer a ~5500 aa protein under monomer mode, using a single 40G card.
sadly, the process halt with error of insufficient video memory, even with chunk size set at only 4 or even smaller.
I noticed that in your demo, if switch to bf16 (or fp16?), sequences < 8000 aa can be inferred with 40G memory.
thus, how do i switch from tf32 to bf16 (or possibly fp16) to save video memory?