Fix evoformer test

tests/test_evoformer.py

@@ -193,10 +193,10 @@ class TestExtraMSAStack(unittest.TestCase):
             ckpt=False,
             inf=inf,
             eps=eps,
-        ).eval()
+        ).eval().cuda()
 
-        m = torch.rand((batch_size, s_t, n_res, c_m))
-        z = torch.rand((batch_size, n_res, n_res, c_z))
+        m = torch.rand((batch_size, s_t, n_res, c_m), device="cuda")
+        z = torch.rand((batch_size, n_res, n_res, c_z), device="cuda")
         msa_mask = torch.randint(
             0,
             2,
@@ -205,6 +205,7 @@ class TestExtraMSAStack(unittest.TestCase):
                 s_t,
                 n_res,
             ),
+            device="cuda",
         )
         pair_mask = torch.randint(
             0,
@@ -214,6 +215,7 @@ class TestExtraMSAStack(unittest.TestCase):
                 n_res,
                 n_res,
             ),
+            device="cuda",
         )
 
         shape_z_before = z.shape