localize

dialoggen/dialoggen_demo.py

@@ -26,8 +26,11 @@ import requests
 from PIL import Image
 from io import BytesIO
 import re
+import devicetorch
 
 
+device = devictorch.get(torch)
+
 def image_parser(image_file, sep=','):
     out = image_file.split(sep)
     return out
@@ -107,7 +110,8 @@ def eval_model(models,
     input_ids = (
         tokenizer_image_token(prompt, models["tokenizer"], IMAGE_TOKEN_INDEX, return_tensors="pt")
         .unsqueeze(0)
-        .cuda()
+        .to(device)
+#        .cuda()
     )
     with torch.inference_mode():
         output_ids = models["model"].generate(