Inference with mixed precision

[soumyadipghosh]

Hello, I am wondering if the inference code can run at mixed precision?

For example, if I modify run_pretrained_openfold.py line 254 from

out = run_model(model, processed_feature_dict, tag, args.output_dir)

to

with torch.cuda.amp.autocast(enabled=True):
        out = run_model(model, processed_feature_dict, tag, args.output_dir)

I see that the plddt score becomes nan. Is this expected behavior?

I am using protein 7ZCV from PDB.

@gahdritz please share your thoughts.

[soumyadipghosh]

Hi @gahdritz @sachinkadyan7 Following up on this to see if you have any thoughts?