use math instead of numpy for prod()

Signed-off-by: Hao Wu <[email protected]>

Author: skyw

emerging_optimizers/utils/modules.py

@@ -13,9 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import math
 from typing import Any, Self
 
-import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -54,7 +54,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self.weight: nn.Parameter[torch.Tensor]
         self.bias: nn.Parameter[torch.Tensor] | None | str
 
-        flat_weight_shape = [self.out_channels, np.prod(self.weight.shape[1:])]
+        flat_weight_shape = [self.out_channels, math.prod(self.weight.shape[1:])]
         if self.bias is not None:
             flat_weight_shape[1] += 1
         flat_weight_buffer = torch.empty(flat_weight_shape, device=self.weight.device, dtype=self.weight.dtype)