Added a --use_cuda_bfloat16 option to the boltz predict command.

src/boltz/main.py

@@ -604,6 +604,11 @@ def cli() -> None:
     is_flag=True,
     help="Whether to not use potentials for steering. Default is False.",
 )
+@click.option(
+    "--use_cuda_bfloat16",
+    is_flag=True,
+    help="Whether to run Boltz model in bfloat16 format when using CUDA instead of float32 to reduce memory use and allow prediction of 40% larger structures.  Prediction accuracy may be reduced.  Has no effect if CUDA is not used.  Default is False.",
+)
 def predict(
     data: str,
     out_dir: str,
@@ -625,6 +630,7 @@ def predict(
     msa_server_url: str = "https://api.colabfold.com",
     msa_pairing_strategy: str = "greedy",
     no_potentials: bool = False,
+    use_cuda_bfloat16: bool = False,
 ) -> None:
     """Run predictions with Boltz-1."""
     # If cpu, write a friendly warning
@@ -761,13 +767,18 @@ def predict(
         precision=32,
     )
 
-    # Compute predictions
-    trainer.predict(
-        model_module,
-        datamodule=data_module,
-        return_predictions=False,
-    )
+    def compute_predictions():
+        trainer.predict(
+            model_module,
+            datamodule=data_module,
+            return_predictions=False,
+        )
 
+    if use_cuda_bfloat16:
+        with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
+            compute_predictions()
+    else:
+        compute_predictions()
 
 if __name__ == "__main__":
     cli()