inpaint and generate are not working with some checkpoints

[adrienchaton]

Hello all,

I could run generate and optimize with the following checkpoint "crossdocked_ca_cond".

However I am encountering different errors depending on checkpoints and methods, would anyone have some fixes please?

1. generate with fullatom checkpoints e.g. "crossdocked_fullatom_cond"

Entropy of n_nodes: H[N] 9.266729354858398
Traceback (most recent call last):
  File " generate_ligands.py", line 51, in <module>
    molecules_batch = model.generate_ligands(
  File "  diffsbdd/src/lightning_modules.py", line 839, in generate_ligands
    self.ddpm.sample_given_pocket(pocket, num_nodes_lig,
  File "  diffsbdd/env/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "  diffsbdd/src/equivariant_diffusion/conditional_model.py", line 525, in sample_given_pocket
    z_lig, xh_pocket = self.sample_p_zs_given_zt(
  File "  diffsbdd/src/equivariant_diffusion/conditional_model.py", line 445, in sample_p_zs_given_zt
    eps_t_lig, _ = self.dynamics(
  File "  diffsbdd/env/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "  diffsbdd/src/equivariant_diffusion/dynamics.py", line 114, in forward
    edges = self.get_edges(mask_atoms, mask_residues, x_atoms, x_residues)
  File "  diffsbdd/src/equivariant_diffusion/dynamics.py", line 185, in get_edges
    edges = torch.stack(torch.where(adj), dim=0)
RuntimeError: nonzero is not supported for tensors with more than INT_MAX elements,   file a support request

2. inpaint with e.g. "crossdocked_ca_cond" ends up with the same error

Entropy of n_nodes: H[N] 7.055830001831055
Traceback (most recent call last):
  File " inpaint.py", line 219, in <module>
    molecules = inpaint_ligand(model, args.pdbfile, args.n_samples,
  File " inpaint.py", line 147, in inpaint_ligand
    xh_lig, xh_pocket, lig_mask, pocket_mask = model.ddpm.inpaint(
  File "  diffsbdd/env/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "  diffsbdd/src/equivariant_diffusion/conditional_model.py", line 632, in inpaint
    z_lig_unknown, xh_pocket = self.sample_p_zs_given_zt(
  File "  diffsbdd/src/equivariant_diffusion/conditional_model.py", line 445, in sample_p_zs_given_zt
    eps_t_lig, _ = self.dynamics(
  File "  diffsbdd/env/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "  diffsbdd/src/equivariant_diffusion/dynamics.py", line 114, in forward
    edges = self.get_edges(mask_atoms, mask_residues, x_atoms, x_residues)
  File "  diffsbdd/src/equivariant_diffusion/dynamics.py", line 185, in get_edges
    edges = torch.stack(torch.where(adj), dim=0)
RuntimeError: nonzero is not supported for tensors with more than INT_MAX elements,   file a support request

3. inpaint with other checkpoints seems to have inconsistencies e.g. "crossdocked_ca_joint"

Entropy of n_nodes: H[N] 7.055830001831055
Traceback (most recent call last):
  File " inpaint.py", line 219, in <module>
    molecules = inpaint_ligand(model, args.pdbfile, args.n_samples,
  File " inpaint.py", line 147, in inpaint_ligand
    xh_lig, xh_pocket, lig_mask, pocket_mask = model.ddpm.inpaint(
  File "  diffsbdd/env/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
TypeError: EnVariationalDiffusion.inpaint() got an unexpected keyword argument 'center'

Also some checkpoints seem less stable than others, generate ran for > 2 hours with "moad_ca_joint" until it crashed due to ValueError: NaN detected in EGNN output, I will try to modify the code to just drop samples when this happen and keep generating new ones at random.

Thanks for any hints!

[adrienchaton]

possibly related to

#28

#55