Allow paramter_overview to work with JAX ShapeDtypeStruct.

PiperOrigin-RevId: 595317017
google · Jan 3, 2024 · 574d4c9 · 574d4c9
1 parent 8a01f21
commit 574d4c9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/clu/parameter_overview.py b/clu/parameter_overview.py
@@ -83,7 +83,7 @@ def _count_parameters(params: _ParamsContainer) -> int:
 def _parameters_size(params: _ParamsContainer) -> int:
   """Returns total size (bytes) for the module or parameter dictionary."""
   params = flatten_dict(params)
-  return sum(v.nbytes for v in params.values())
+  return sum(np.prod(v.shape) * v.dtype.itemsize for v in params.values())
 
 
 def count_parameters(params: _ParamsContainer) -> int: