fix CUDA with MDS

lucidrains · Feb 3, 2021 · 8fc6d12 · 8fc6d12
1 parent 5f9355b
commit 8fc6d12
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/alphafold2_pytorch/utils.py b/alphafold2_pytorch/utils.py
@@ -262,6 +262,8 @@ def mds_torch(pre_dist_mat, weights=None, iters=10, tol=1e-5, verbose=2):
         * best_3d_coords: (3 x N)
         * historic_stress 
     """
+    device = pre_dist_mat.device
+
     if weights is None:
         weights = torch.ones_like(pre_dist_mat)
     # batched MDS
@@ -271,8 +273,8 @@ def mds_torch(pre_dist_mat, weights=None, iters=10, tol=1e-5, verbose=2):
     batch, N, _ = pre_dist_mat.shape
     his = []
     # init random coords
-    best_stress = float("Inf") * torch.ones(batch) 
-    best_3d_coords = 2*torch.rand(batch, N, 3) - 1
+    best_stress = float("Inf") * torch.ones(batch, device = device)
+    best_3d_coords = 2*torch.rand(batch, N, 3, device = device) - 1
     # iterative updates:
     for i in range(iters):
         # compute distance matrix of coords and stress

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'alphafold2-pytorch',
   packages = find_packages(),
-  version = '0.0.6',
+  version = '0.0.7',
   license='MIT',
   description = 'AlphaFold2 - Pytorch',
   author = 'Phil Wang, Eric Alcaide',