Merge branch 'main' into bug/onnx_params

onnx · Mar 1, 2024 · 1aa3e7d · 1aa3e7d
2 parents b4aa3d7 + 293e903
commit 1aa3e7d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/turnkeyml/analyze/model.py b/src/turnkeyml/analyze/model.py
@@ -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
             )