From 1db9d9ba61fc79f94e6236f2d63416623d928716 Mon Sep 17 00:00:00 2001 From: onesome Date: Tue, 25 Apr 2023 18:46:54 -0500 Subject: [PATCH] Lazyload: Whoops --- modeling/lazy_loader.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modeling/lazy_loader.py b/modeling/lazy_loader.py index 36a32cea..3dee5bae 100644 --- a/modeling/lazy_loader.py +++ b/modeling/lazy_loader.py @@ -389,7 +389,7 @@ def safetensors_load_tensor_independently( return f.get_tensor(tensor_key) -def patch_safetensors(): +def patch_safetensors(callback): # Safetensors load patch import transformers @@ -494,7 +494,7 @@ def use_lazy_load( torch.load = torch_load if HAS_SAFETENSORS: - patch_safetensors() + patch_safetensors(callback) if dematerialized_modules: if use_accelerate_init_empty_weights: