fix device

app.py

@@ -237,7 +237,6 @@ def buildmodel(**kwargs):
     seed = kwargs.get('seed', None)
     input_model = kwargs.get('input_model', None)
     quantizer_version = kwargs.get('quantizer_version', 'v4')
-    device = "cuda"
     
     set_seed(seed)
     # old_tokenizer = AutoTokenizer.from_pretrained(input_model, trust_remote_code=True)
@@ -289,7 +288,7 @@ def buildmodel(**kwargs):
         subfolder="checkpoint-26000",      # 加载检查点目录
         config=model_args,
         # cache_dir="/openseg_blob/v-yanbin/GradioDemo/cache_dir",
-    )
+    ).to("cuda")
     # model = CrelloModel(config=model_args)
     
     tokenizer.add_special_tokens({"mask_token": "<mask>"}) 
@@ -327,7 +326,7 @@ def construction_layout():
     # print('after token embeddings to match the tokenizer', 129423)
     
     print("before .to(device)")
-    model = model.to(device)
+    model = model.to("cuda")
     print("after .to(device)")
     model = model.bfloat16()
     model.eval()
@@ -339,8 +338,8 @@ def evaluate_v1(inputs, model, quantizer, tokenizer, width, height, device, do_s
     json_example = inputs
     input_intension = '{"wholecaption":"' + json_example["wholecaption"] + '","layout":[{"layer":'
     inputs = tokenizer(
-            input_intension, return_tensors="pt"
-        ).to(device)
+        input_intension, return_tensors="pt"
+    ).to("cuda")
     
     stopping_criteria = StoppingCriteriaList()
     stopping_criteria.append(StopAtSpecificTokenCriteria(token_id_list=[128000]))

app_test.py

@@ -237,7 +237,6 @@ def buildmodel(**kwargs):
     seed = kwargs.get('seed', None)
     input_model = kwargs.get('input_model', None)
     quantizer_version = kwargs.get('quantizer_version', 'v4')
-    device = "cuda"
     
     set_seed(seed)
     # old_tokenizer = AutoTokenizer.from_pretrained(input_model, trust_remote_code=True)
@@ -289,7 +288,7 @@ def buildmodel(**kwargs):
         subfolder="checkpoint-26000",      # 加载检查点目录
         config=model_args,
         cache_dir="/openseg_blob/v-yanbin/GradioDemo/cache_dir",
-    )
+    ).to("cuda")
     # model = CrelloModel(config=model_args)
     
     tokenizer.add_special_tokens({"mask_token": "<mask>"}) 
@@ -327,7 +326,7 @@ def construction_layout():
     # print('after token embeddings to match the tokenizer', 129423)
     
     print("before .to(device)")
-    model = model.to(device)
+    model = model.to("cuda")
     print("after .to(device)")
     model = model.bfloat16()
     model.eval()
@@ -339,8 +338,8 @@ def evaluate_v1(inputs, model, quantizer, tokenizer, width, height, device, do_s
     json_example = inputs
     input_intension = '{"wholecaption":"' + json_example["wholecaption"] + '","layout":[{"layer":'
     inputs = tokenizer(
-            input_intension, return_tensors="pt"
-        ).to(device)
+        input_intension, return_tensors="pt"
+    ).to("cuda")
     
     stopping_criteria = StoppingCriteriaList()
     stopping_criteria.append(StopAtSpecificTokenCriteria(token_id_list=[128000]))