2

app.py

@@ -342,7 +342,7 @@ def evaluate_v1(inputs, model, quantizer, tokenizer, width, height, device, do_s
     input_intension = '{"wholecaption":"' + json_example["wholecaption"] + '","layout":[{"layer":'
     inputs = tokenizer(
         input_intension, return_tensors="pt"
-    ).to("cuda")
+    ).to(model.lm.device)
     
     stopping_criteria = StoppingCriteriaList()
     stopping_criteria.append(StopAtSpecificTokenCriteria(token_id_list=[128000]))

app_test.py

@@ -342,7 +342,7 @@ def evaluate_v1(inputs, model, quantizer, tokenizer, width, height, device, do_s
     input_intension = '{"wholecaption":"' + json_example["wholecaption"] + '","layout":[{"layer":'
     inputs = tokenizer(
         input_intension, return_tensors="pt"
-    ).to("cuda")
+    ).to(model.lm.device)
     
     stopping_criteria = StoppingCriteriaList()
     stopping_criteria.append(StopAtSpecificTokenCriteria(token_id_list=[128000]))