From 5539e387c08a54c77b9cb5e92faf7d878bff40ad Mon Sep 17 00:00:00 2001
From: blackhao <13851610112@163.com>
Date: Sat, 23 Aug 2025 15:03:34 -0500
Subject: maybe

---
 Group-Entropy-Equalization/train.py | 1 +
 1 file changed, 1 insertion(+)

(limited to 'Group-Entropy-Equalization/train.py')

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")
-- 
cgit v1.2.3