_LinAlgErr when running inference

[ingcoder]

Hi, I'm using the latest version of DiffDock (1.1.2) on an M1 architecture and encountering an error during inference on a PDB file with two protein chains. The issue seems related to how the 'rot_score' tensor is generated, which occasionally includes NaN values. Sometimes the code works, and other times it doesn't.

Has anyone else experienced this problem and knows how to resolve it?

The error I believe originates in line 115 in sampling.py.
tr_score, rot_score, tor_score = model(mod_complex_graph_batch)[:3]

rot_score tensor

([[-0.0200, 0.0089, 0.0849], [-0.0444, -0.0414, -0.0024], [ 0.1594, -0.0008, 0.0887], [ 0.0368, -0.1539, -0.1124], [ 0.1238, 0.0957, 0.1374], [-0.0850, -0.0477, 0.0188], [-0.0089, -0.0112, 0.0326], [ nan, nan, nan], [ 0.0850, -0.1109, 0.1254], [ 0.0040, -0.0470, -0.0599]])

The rot_score is used in the rot_pertub calculation which then also produces nan.
rot_perturb = (rot_score * dt_rot * rot_g ** 2 + rot_g * np.sqrt(dt_rot) * rot_z)

Resulting in runtime error:

Traceback (most recent call last):
  File "/DiffDock/inference.py", line 262, in main
    data_list, confidence = sampling(data_list=data_list, model=model,
  File "/DiffDock/utils/sampling.py", line 178, in sampling
    modify_conformer_batch(complex_graph_batch['ligand'].pos, complex_graph_batch, tr_perturb, rot_perturb,
  File "/DiffDock/utils/diffusion_utils.py", line 73, in modify_conformer_batch
    R, t = rigid_transform_Kabsch_3D_torch_batch(flexible_new_pos, rigid_new_pos)
  File "/DiffDock/utils/geometry.py", line 267, in rigid_transform_Kabsch_3D_torch_batch
    U, S, Vt = torch.linalg.svd(H)
torch._C._LinAlgError: linalg.svd: (Batch element 0): The algorithm failed to converge because the input matrix contained non-finite values.