From 103ef760dac64295823e3a5f5890344a778a0dd7 Mon Sep 17 00:00:00 2001 From: Alexander Borzunov Date: Wed, 24 Jul 2024 04:13:40 -0700 Subject: [PATCH] Materialize buffers in get_block_size() (#600) --- src/petals/server/block_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/petals/server/block_utils.py b/src/petals/server/block_utils.py index 63ba68763..3ba739744 100644 --- a/src/petals/server/block_utils.py +++ b/src/petals/server/block_utils.py @@ -32,7 +32,7 @@ def get_block_size( dtype is not None and quant_type is not None ), 'get_block_size(..., location="memory") requires to specify dtype and quant_type for calculations' - with init_empty_weights(include_buffers=True): + with init_empty_weights(include_buffers=False): block = get_model_block(config) n_params = sum(param.numel() for param in block.parameters())