Auto-calculate total model bytes on model descriptor

chaiNNer-org · May 17, 2024 · 25661d8 · 25661d8
1 parent 65ebd7e
commit 25661d8
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/libs/spandrel/spandrel/__helpers/model_descriptor.py b/libs/spandrel/spandrel/__helpers/model_descriptor.py
@@ -214,6 +214,11 @@ def __init__(
         on how to best use the model.
         """
 
+        self.model_bytes: int = sum(p.numel() * 4 for p in model.parameters())
+        """
+        The number of bytes the model uses in fp32.
+        """
+
         self.model.load_state_dict(state_dict)  # type: ignore
 
     @property