Number of parameters must be an integer

Signed-off-by: Joao Lucas de Sousa Almeida <[email protected]>
IBM · May 24, 2023 · 07c7419 · 07c7419
1 parent dc7c1a5
commit 07c7419
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/simulai/templates/_pytorch_network.py b/simulai/templates/_pytorch_network.py
@@ -58,8 +58,8 @@ def weights_l1(self) -> torch.Tensor:
     def n_parameters(self):
 
         try:
-            return sum([np.prod(tuple(param.shape)) for param in
-                    self.parameters()])
+            return int(sum([np.prod(tuple(param.shape)) for param in
+                    self.parameters()]))
 
         except Exception: