fix device

1049451037 · web-flow · commit e314fb9c4e77 · 2024-08-23T15:45:06.000+08:00
diff --git a/finetune_visualglm.py b/finetune_visualglm.py
@@ -178,6 +178,7 @@ def create_dataset_function(path, args):
     model, args = FineTuneVisualGLMModel.from_pretrained(model_type, args)
     if torch.cuda.is_available():
         model = model.to('cuda')
+        args.device = 'cuda'
     tokenizer = get_tokenizer(args)
     label_pad_token_id = -100 if args.ignore_pad_token_for_loss else tokenizer.pad_token_id
     def data_collator(examples):
