Skip to content

Commit

Permalink
Avoid parameter overflow when counting parameters of very large models (
Browse files Browse the repository at this point in the history
  • Loading branch information
danielholanda authored Mar 1, 2024
1 parent ef90ccf commit 293e903
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion src/turnkeyml/analyze/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ def count_parameters(model: torch.nn.Module, model_type: build.ModelType) -> int
onnx_model = onnx.load(model)
return int(
sum(
np.prod(tensor.dims)
np.prod(tensor.dims, dtype=np.int64)
for tensor in onnx_model.graph.initializer
if tensor.name not in onnx_model.graph.input
)
Expand Down

0 comments on commit 293e903

Please sign in to comment.