SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Zero shot distillation script cuda patch (#10284)

Browse Source
This commit is contained in:
Joe Davison
2021-02-19 14:06:57 -05:00
committed by GitHub
parent f1299f5038
commit cbadb5243c
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
examples/research_projects/zero-shot-distillation/distill_classifier.py
View File

@@ -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)