fixed document (#13414)

41cd52a7 · Mohan Zhang · GitHub · 330d83fd · 41cd52a7
Unverified Commit 41cd52a7 authored Sep 08, 2021 by Mohan Zhang Committed by GitHub Sep 08, 2021
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

docs/source/main_classes/trainer.rst docs/source/main_classes/trainer.rst +2 -2

No files found.
--- a/docs/source/main_classes/trainer.rst
+++ b/docs/source/main_classes/trainer.rst
@@ -64,9 +64,9 @@ classification:

    class MultilabelTrainer(Trainer):
        def compute_loss(self, model, inputs, return_outputs=False):
-            labels = inputs.pop("labels")
+            labels = inputs.get("labels")
            outputs = model(**inputs)
-            logits = outputs.logits
+            logits = outputs.get('logits')
            loss_fct = nn.BCEWithLogitsLoss()
            loss = loss_fct(logits.view(-1, self.model.config.num_labels),
                            labels.float().view(-1, self.model.config.num_labels))