Open oliverdutton opened 2 months ago
Matches the casting in multimer for monomers
Runtime speedup of ~1.8x and reduces peak memory usage ~2x (for 759 residue protein, model_5)
Here's a colab link showing runtime improvement and no significant change in prediction output by visual inspection
No effect on default usage as bfloat16 turned off.
I've checked a profile on tensorboard and there's no float32 gemms with significant runtime.
Matches the casting in multimer for monomers
Runtime speedup of ~1.8x and reduces peak memory usage ~2x (for 759 residue protein, model_5)
Here's a colab link showing runtime improvement and no significant change in prediction output by visual inspection
No effect on default usage as bfloat16 turned off.
I've checked a profile on tensorboard and there's no float32 gemms with significant runtime.