See comment here.
We should update this so that newer versions of PyTorch do not complain. Could also reimplement the inference to use a model zoo.