From 03ef918dd8efb22d4b0561b281420795ec028c89 Mon Sep 17 00:00:00 2001 From: D-X-Y <280835372@qq.com> Date: Mon, 8 Mar 2021 08:24:23 +0000 Subject: [PATCH] Fix bugs in count_parameters --- qlib/contrib/model/pytorch_utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/qlib/contrib/model/pytorch_utils.py b/qlib/contrib/model/pytorch_utils.py index f2457c9a6..c0d483db3 100644 --- a/qlib/contrib/model/pytorch_utils.py +++ b/qlib/contrib/model/pytorch_utils.py @@ -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":