Add fix for huggingface/transformers#8237 (comment)

nelson-liu · Apr 6, 2021 · 10daafb · 10daafb
1 parent eeb7877
commit 10daafb
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/transformers/trainer_pt_utils.py b/src/transformers/trainer_pt_utils.py
@@ -34,7 +34,7 @@
 if is_torch_tpu_available():
     import torch_xla.core.xla_model as xm
 
-if version.parse(torch.__version__) <= version.parse("1.4.1"):
+if version.parse(torch.__version__) <= version.parse("1.4.1") or version.parse(torch.__version__) > version.parse("1.7.0"):
     SAVE_STATE_WARNING = ""
 else:
     from torch.optim.lr_scheduler import SAVE_STATE_WARNING