Zero shot distillation script cuda patch (#10284)

huggingface · Feb 19, 2021 · cbadb52 · cbadb52
1 parent f1299f5
commit cbadb52
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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)