Zero shot distillation script cuda patch (#10284)

cbadb524 · Joe Davison · GitHub · f1299f50 · cbadb524
Unverified Commit cbadb524 authored Feb 19, 2021 by Joe Davison Committed by GitHub Feb 19, 2021
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

examples/research_projects/zero-shot-distillation/distill_classifier.py ...rch_projects/zero-shot-distillation/distill_classifier.py +1 -1

No files found.
--- a/examples/research_projects/zero-shot-distillation/distill_classifier.py
+++ b/examples/research_projects/zero-shot-distillation/distill_classifier.py
@@ -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)