now check the batch_size in compute_tm function

openfold/utils/loss.py

@@ -698,7 +698,8 @@ def compute_tm(
     if interface and (asym_id is not None):
         if len(asym_id.shape)>1:
             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)
     
     predicted_tm_term *= pair_mask