bugfix: removing nans from alphafold-multimer+templates

google-deepmind#513

alphafold/model/utils.py

@@ -34,7 +34,7 @@ def final_init(config):
 
 def batched_gather(params, indices, axis=0, batch_dims=0):
   """Implements a JAX equivalent of `tf.gather` with `axis` and `batch_dims`."""
-  take_fn = lambda p, i: jnp.take(p, i, axis=axis)
+  take_fn = lambda p, i: jnp.take(p, i, axis=axis, mode="clip")
   for _ in range(batch_dims):
     take_fn = jax.vmap(take_fn)
   return take_fn(params, indices)