fix autocast for older pytorch

14cc50d0 · Stas Bekman · 4c0dd199 · 14cc50d0
Commit 14cc50d0 authored Dec 01, 2021 by Stas Bekman
Hide whitespace changes
Inline Side-by-side

Showing with 18 additions and 6 deletions

src/transformers/trainer.py src/transformers/trainer.py +18 -6

No files found.
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -1857,8 +1857,12 @@ class Trainer:
            return loss_mb.reduce_mean().detach().to(self.args.device)
        if self.use_amp:
-            with autocast(dtype=self.amp_dtype):
+            if version.parse(torch.__version__) >= version.parse("1.10"):
-                loss = self.compute_loss(model, inputs)
+                with autocast(dtype=self.amp_dtype):
+                    loss = self.compute_loss(model, inputs)
+            else:
+                with autocast():
+                    loss = self.compute_loss(model, inputs)
        else:
            loss = self.compute_loss(model, inputs)
@@ -2501,8 +2505,12 @@ class Trainer:
            else:
                if has_labels:
                    if self.use_amp:
-                        with autocast(dtype=self.amp_dtype):
+                        if version.parse(torch.__version__) >= version.parse("1.10"):
-                            loss, outputs = self.compute_loss(model, inputs, return_outputs=True)
+                            with autocast(dtype=self.amp_dtype):
+                                loss, outputs = self.compute_loss(model, inputs, return_outputs=True)
+                        else:
+                            with autocast():
+                                loss, outputs = self.compute_loss(model, inputs, return_outputs=True)
                    else:
                        loss, outputs = self.compute_loss(model, inputs, return_outputs=True)
                    loss = loss.mean().detach()
@@ -2514,8 +2522,12 @@ class Trainer:
                else:
                    loss = None
                    if self.use_amp:
-                        with autocast(dtype=self.amp_dtype):
+                        if version.parse(torch.__version__) >= version.parse("1.10"):
-                            outputs = model(**inputs)
+                            with autocast(dtype=self.amp_dtype):
+                                outputs = model(**inputs)
+                        else:
+                            with autocast():
+                                outputs = model(**inputs)
                    else:
                        outputs = model(**inputs)
                    if isinstance(outputs, dict):