Investigate and potentially implement `torch.compile()` in `train.py`

There is a TODO comment in the `train.py` script suggesting the potential use of `torch.compile()`. This issue is to investigate whether `torch.compile()` can be beneficial for the project and, if so, to plan its implementation.

Related PR: https://github.com/talmolab/biogtr/pull/24
Comment URL: https://github.com/talmolab/biogtr/pull/24#discussion_r1575428228