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?