Training for vector and tensor variables

[chuyuanchen]

Dear developers, currently tensor information in Allegro is abandoned in the end and a scalar energy is produced. However, I think the equivariant latent space can be useful to train vector and tensor variables which have rotation symmetry? For example, the atom wise born effective charge (a rank 2 tensor), or the polarization vector, which describe the linear response of the systems (force and energy) to external electric field and would allow dynamics simulation under external field. Is there any plan to implement such feature to train for vectors and higher tensors? I think all is needed is to pick up the correct irreps in the final layer (e.g., instead of "0e", output "0e+1o" for vector) and modify the loss function, does it make sense to you? Thank you very much.