Update main.py

main.py

@@ -165,7 +165,7 @@ def model_init():
     # args = set_args()
     acuda = torch.cuda.is_available() and not args.no_cuda
     device = 'cuda' if acuda else 'cpu'
-    os.environ["CUDA_VISIBLE_DEVICES"] = args.device
+    os.environ["CUDA_VISIBLE_DEVICES"] = device
     tokenizer = BertTokenizerFast(vocab_file='/app/bert-base-zh/vocab.txt', sep_token="[SEP]", pad_token="[PAD]", cls_token="[CLS]")
     # # tokenizer = BertTokenizer(vocab_file=args.voca_path)
     model = Word_BERT()