restructure embeddings to avoid issues with torch jit

[stefaanhessmann]

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.

This is still a draft and not tested!