Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
A couple of PRs ago, we decided to include
dtype
anddevice
as explicit and obligatory parameters for bothcalculators
andpotentials
.Unfortunately, after thorough consideration of how typical pipelines are built, I concluded that we should abandon this design choice.
The main reason is that, in most cases, when working with an NN
model,
the preferred strategy is to first initialize themodel
and then move it to the desired device usingmodel.to(device)
.Since
torch-pme
is designed to be an internal part of themodel
, this creates a conflict. We initializedtype
anddevice
once, but when we later move themodel
to a differentdevice
, it undermines our prior device-checking logic.Luckily, since our entire pipeline is either a
torch.nn.Module
or its subclass, we can integrate it smoothly with models that change theirdevice
anddtype
. The key idea is to thoroughly rewrite the pipeline so that all newly created tensors during calculations are registered as buffers usingself.register_buffer
.This PR aims to achieve exactly that.
📚 Documentation preview 📚: https://torch-pme--166.org.readthedocs.build/en/166/