diff options
Diffstat (limited to 'Group-Entropy-Equalization/train.py')
| -rw-r--r-- | Group-Entropy-Equalization/train.py | 1 |
1 files changed, 1 insertions, 0 deletions
diff --git a/Group-Entropy-Equalization/train.py b/Group-Entropy-Equalization/train.py index 9ca1d4a..352e916 100644 --- a/Group-Entropy-Equalization/train.py +++ b/Group-Entropy-Equalization/train.py @@ -18,6 +18,7 @@ from accelerate import Accelerator, DeepSpeedPlugin from accelerate.utils import set_seed from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM import json +import math os.environ.setdefault("NCCL_TIMEOUT", "2700") |