diff --git a/lora/convert.py b/lora/convert.py index 153e056f..2ce247a3 100644 --- a/lora/convert.py +++ b/lora/convert.py @@ -32,7 +32,12 @@ def map_torch_to_mlx(key, value): elif "rope" in key: return None, None - return key, value.numpy() + return ( + key, + value.numpy() + if value.dtype != torch.bfloat16 + else value.to(torch.float32).numpy(), + ) if __name__ == "__main__":