"launch/git@developer.sourcefind.cn:OpenDAS/dynamo.git" did not exist on "86aff237ab4c098587c89ce9982f913ed32a87a8"
Commit a420160f authored by Geoffrey Yu's avatar Geoffrey Yu
Browse files

now check the batch_size in compute_tm function

parent 5e1af36e
...@@ -698,7 +698,8 @@ def compute_tm( ...@@ -698,7 +698,8 @@ def compute_tm(
if interface and (asym_id is not None): if interface and (asym_id is not None):
if len(asym_id.shape)>1: if len(asym_id.shape)>1:
assert len(asym_id.shape)<=2 assert len(asym_id.shape)<=2
pair_mask = residue_weights.new_ones((1,n, n), dtype=torch.int32) batch_size = asym_id.shape[0]
pair_mask = residue_weights.new_ones((batch_size,n, n), dtype=torch.int32)
pair_mask *= (asym_id[..., None] != asym_id[..., None, :]).to(dtype=pair_mask.dtype) pair_mask *= (asym_id[..., None] != asym_id[..., None, :]).to(dtype=pair_mask.dtype)
predicted_tm_term *= pair_mask predicted_tm_term *= pair_mask
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment