From be2a5cf3a666de5824fc1304eef369ce296f2321 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 14 Mar 2025 02:07:13 -0700 Subject: [PATCH] Update _utils.py --- unsloth/models/_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index af423d560f..b062bc2d49 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -184,7 +184,7 @@ except: # Patch get_model_param_count to record correct 4bit / 8bit from transformers.trainer_pt_utils import is_deepspeed_zero3_enabled -def get_model_param_count(model, trainable_only=False): +def get_model_param_count(model, trainable_only = False): """ Calculate model's total param count. If trainable_only is True then count only those requiring grads """ @@ -208,6 +208,8 @@ def get_model_param_count(model, trainable_only=False): pass import transformers.trainer_pt_utils transformers.trainer_pt_utils.get_model_param_count = get_model_param_count +import transformers.trainer +transformers.trainer.get_model_param_count = get_model_param_count # ============================================= # =============================================