saymrwulf/transformers
1
0
Fork 0
mirror of https://github.com/saymrwulf/transformers.git synced 2026-05-14 20:58:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

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
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show stats 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_config = model.config
if not no_cuda and torch.cuda.is_available():
model = nn.DataParallel(model)
model = nn.DataParallel(model.cuda())
batch_size *= len(model.device_ids)
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=use_fast_tokenizer)