turn off mixed precision in weighted rigid align due to svd

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -10,6 +10,7 @@
 import torch
 from torch import nn
 from torch import Tensor
+from torch.amp import autocast
 import torch.nn.functional as F
 from torch.utils.checkpoint import checkpoint, checkpoint_sequential
 
@@ -2979,6 +2980,7 @@ class WeightedRigidAlign(Module):
     """ Algorithm 28 """
 
     @typecheck
+    @autocast('cuda', enabled = False)
     def forward(
         self,
         pred_coords: Float['b n 3'],       # predicted coordinates

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.4.5"
+version = "0.4.6"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" },