I restructured the embeddings to avoid problems with torch jit.
Nuclear and electronic embeddings are separated, because electronic embeddings currently require (embedding, inputs) as arguments in the forward pass and nuclear embeddings only require atomic_numbers in the forward pass.
With this, we can further use torch.nn.Embedding as a default value and electronic embeddings can be added as nn.ModuleList.
I restructured the embeddings to avoid problems with torch jit.
Nuclear and electronic embeddings are separated, because electronic embeddings currently require
(embedding, inputs)as arguments in the forward pass and nuclear embeddings only requireatomic_numbersin the forward pass. With this, we can further usetorch.nn.Embeddingas a default value and electronic embeddings can be added asnn.ModuleList.This is still a draft and not tested!