Skip to content

Conversation

@zhubonan
Copy link
Contributor

Add a batch relaxer which combines multiple atoms as a single batch to the underlying model. This can give speed-ups when lots of smaller structures needs to be relaxed using GPU (typical while using the foundation models for high-throughput studies).

The current code is working, but a few useful functionalities can still be added:

  • Support for ensemble models
  • Support other types of MACE like EnergyDipoleMACE?
  • Since we already run the optimiser step-by-step, it does no harm to add some detection for pathological structures (exploded/collapsed) which can make the main worker process to stuck.
  • Automatically determine the best batch size by collecting the timings of the compute method.

@zhubonan zhubonan marked this pull request as draft March 10, 2025 15:40
@ilyes319
Copy link
Contributor

Hey looks great, ping me when you think it is ready to merge.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants