Skip to content

GitLab

"git@developer.sourcefind.cn:chenpangpang/transformers.git" did not exist on "b242d0f297aa87a0c8d99657a53691ece2dfe492"
Unverified Commit cbadb524 authored by Joe Davison's avatar Joe Davison Committed by GitHub
Browse files

Zero shot distillation script cuda patch (#10284)

parent f1299f50
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
examples/research_projects/zero-shot-distillation/distill_classifier.py
View file @ cbadb524
...@@ -174,7 +174,7 @@ def get_teacher_predictions( ...@@ -174,7 +174,7 @@ def get_teacher_predictions(
model = AutoModelForSequenceClassification.from_pretrained(model_path) model = AutoModelForSequenceClassification.from_pretrained(model_path)
model_config = model.config model_config = model.config
if not no_cuda and torch.cuda.is_available(): if not no_cuda and torch.cuda.is_available():
model = nn.DataParallel(model) model = nn.DataParallel(model.cuda())
batch_size *= len(model.device_ids) batch_size *= len(model.device_ids)
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=use_fast_tokenizer) tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=use_fast_tokenizer)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment