From 8c8046cce51c65a596290179e65715b640062a65 Mon Sep 17 00:00:00 2001 From: Taksh Date: Fri, 17 Jul 2026 11:04:58 +0300 Subject: [PATCH] fix: honor dtype/device in load_state_dict_hf Early return skipped the conversion block so dtype= and device= were ignored. --- mamba_ssm/utils/hf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mamba_ssm/utils/hf.py b/mamba_ssm/utils/hf.py index 0d7555acd..778bd1851 100644 --- a/mamba_ssm/utils/hf.py +++ b/mamba_ssm/utils/hf.py @@ -15,7 +15,7 @@ def load_state_dict_hf(model_name, device=None, dtype=None): # If not fp32, then we don't want to load directly to the GPU mapped_device = "cpu" if dtype not in [torch.float32, None] else device resolved_archive_file = cached_file(model_name, WEIGHTS_NAME, _raise_exceptions_for_missing_entries=False) - return torch.load(resolved_archive_file, map_location=mapped_device) + state_dict = torch.load(resolved_archive_file, map_location=mapped_device) # Convert dtype before moving to GPU to save memory if dtype is not None: state_dict = {k: v.to(dtype=dtype) for k, v in state_dict.items()}