Adds float wrapper to to vectors in TestExtraMSAStack in

test_evoformer.py

Adds float wrapper to to vectors in TestExtraMSAStack in
test_evoformer.py
48668ca3 · Jennifer Wei · 2134cc09 · 48668ca3
Commit 48668ca3 authored Sep 18, 2023 by Jennifer Wei
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

tests/test_evoformer.py tests/test_evoformer.py +2 -2

No files found.
--- a/tests/test_evoformer.py
+++ b/tests/test_evoformer.py
@@ -206,7 +206,7 @@ class TestExtraMSAStack(unittest.TestCase):
                n_res,
            ),
            device="cuda",
-        )
+        ).float()
        pair_mask = torch.randint(
            0,
            2,
@@ -216,7 +216,7 @@ class TestExtraMSAStack(unittest.TestCase):
                n_res,
            ),
            device="cuda",
-        )
+        ).float()

        shape_z_before = z.shape