Support for newer PyTorch versions

[ant0nsc]

This is related to #3.

We are running A6000 GPUs that are not compatible with PyTorch 1.6. Also, our whole codebase is on PyTorch 2.0. We are unable to build any Conda environment that is compatible with the package here. This makes it really cumbersome to use the metrics provided by this repo: we always have to compute metrics in separate processes, and on separate older machines - this is super awkward and time consuming

At the same time, the metrics provided by the package here are super valuable, and we would love to use them more, and in better ways than we can currently

A release of this package that runs with current PyTorch, or at least with a set of more recent PyTorch versions, would be super helpful for our work.
@KathyFeiyang you seem to be most active here on the repo, is that something your team could provide?