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