From bb965d8e870e0e642e0dca079bb91bd5e41e15a3 Mon Sep 17 00:00:00 2001 From: wineandchord Date: Thu, 13 Mar 2025 22:27:50 +0800 Subject: [PATCH] fix type annotation for ALL_ATTENTION_FUNCTIONS (#36690) Corrects the type annotation to match actual usage. The variable was typed as Dict[str, Dict[str, Callable]] but is actually used as Dict[str, Callable] where keys are attention mechanism names and values are the corresponding attention functions directly. This change makes the type annotation consistent with how the dictionary is used in the codebase. --- src/transformers/modeling_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 77f842aa5f..fc94ec093d 100755 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -5835,7 +5835,7 @@ def get_disk_only_shard_files(device_map, weight_map): return [fname for fname, devices in files_content.items() if set(devices) == {"disk"}] -ALL_ATTENTION_FUNCTIONS: Dict[str, Dict[str, Callable]] = {} +ALL_ATTENTION_FUNCTIONS: Dict[str, Callable] = {} ALL_ATTENTION_FUNCTIONS.update( {