Materialize buffers in get_block_size() (#600)

bigscience-workshop · Jul 24, 2024 · 103ef76 · 103ef76
1 parent 10f7525
commit 103ef76
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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())