fix type after model.float() (dptech-corp#92)

hypnopump · Jan 12, 2023 · 413c78b · 413c78b
1 parent ab4b89a
commit 413c78b
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/unifold/model.py b/unifold/model.py
@@ -37,6 +37,10 @@ def half(self):
     def bfloat16(self):
         self.model = self.model.bfloat16()
         return self
+
+    def float(self):
+        self.model = self.model.float()
+        return self
 
     @classmethod
     def build_model(cls, args, task):

diff --git a/unifold/modules/alphafold.py b/unifold/modules/alphafold.py
@@ -117,6 +117,11 @@ def bfloat16(self):
             self.__make_input_float__()
         self.dtype = torch.bfloat16
         return self
+
+    def float(self):
+        super().float()
+        self.dtype = torch.float
+        return self
 
     def alphafold_original_mode(self):
         def set_alphafold_original_mode(module):