Make infer general, so it runs on non cuda devices
#6
by
aamirshakir
- opened
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.