predict.py with using gpu

saahiluppal / catr

Image Captioning Using Transformer

Apache License 2.0

254 stars 53 forks source link

Open batooooo opened 2 years ago

batooooo commented 2 years ago

Hi, I want to get results using gpu, what should i do?

saahiluppal commented 2 years ago

Hi, move model to cuda model.to(device) move input image to cuda input.to(device)