Just a few small suggestions here. I added print(f"CHGNet will run on {self.device}") as I find it helpful if the model reports what device it will be running on so I can see it in the slurm logs later. Helps me make sure I'm using GPU as expected.
Also added some literal type hints for targets: Literal["ef", "efs", "efsm"]. Modern IDEs will surface this information so you don't have to read the code or the docs to know the options.
Just a few small suggestions here. I added
print(f"CHGNet will run on {self.device}")as I find it helpful if the model reports what device it will be running on so I can see it in theslurmlogs later. Helps me make sure I'm using GPU as expected.Also added some literal type hints for
targets: Literal["ef", "efs", "efsm"]. Modern IDEs will surface this information so you don't have to read the code or the docs to know the options.