Make infer general, so it runs on non cuda devices

#6

Infer hard codes .cuda(), which does not allow it to run on CPU or other devices, instead use self.device so it can run on CPU, MPS etc.

Ah, linting, sorry let me update it.

Ready to merge
This branch is ready to get merged automatically.

Sign up or log in to comment