liuyu/qlib
1
0
Fork 0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-07 13:00:58 +08:00
Code Activity

Fix bugs in count_parameters

Browse Source
This commit is contained in:
D-X-Y
2021-03-08 08:24:23 +00:00
parent ca48345b29
commit 03ef918dd8
1 changed files with 1 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
qlib/contrib/model/pytorch_utils.py
View File

@@ -3,14 +3,13 @@
import torch.nn as nn
def count_parameters(models_or_parameters, unit="mb"):
if isinstance(models_or_parameters, nn.Module):
counts = sum(v.numel() for v in models_or_parameters.parameters())
elif isinstance(models_or_parameters, nn.Parameter):
counts = models_or_parameters.numel()
elif isinstance(models_or_parameters, (list, tuple)):
return sum(count_parameters(x, unit) for x in model_or_parameters)
return sum(count_parameters(x, unit) for x in models_or_parameters)
else:
counts = sum(v.numel() for v in models_or_parameters)
if unit.lower() == "mb":