Device bug fix

custom_got/modeling_GOT.py

@@ -20,9 +20,9 @@ DEFAULT_IM_START_TOKEN = '<img>'
 DEFAULT_IM_END_TOKEN = '</img>'
 cuda_is_available = torch.cuda.is_available()
 if cuda_is_available:
-    device = torch.device('cuda')
+    device_type = torch.device('cuda').type
 else:
-    device = torch.device('cpu')
+    device_type = torch.device('cpu').type
 
 from enum import auto, Enum
 class SeparatorStyle(Enum):
@@ -574,7 +574,7 @@ class GOTQwenForCausalLM(Qwen2ForCausalLM):
         streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 
         if stream_flag:
-            with torch.autocast(device, dtype=torch.bfloat16):
+            with torch.autocast(device_type, dtype=torch.bfloat16):
                 output_ids = self.generate(
                     input_ids,
                     images=[image_tensor_1.unsqueeze(0).half().cuda()],
@@ -586,7 +586,7 @@ class GOTQwenForCausalLM(Qwen2ForCausalLM):
                     stopping_criteria=[stopping_criteria]
                     )
         else:
-            with torch.autocast(device, dtype=torch.bfloat16):
+            with torch.autocast(device_type, dtype=torch.bfloat16):
                 output_ids = self.generate(
                     input_ids,
                     images=[image_tensor_1.unsqueeze(0).half().cuda()] if cuda_is_available else [image_tensor_1.unsqueeze(0).half()],