From 85b86f3667d5d228336337d14335fc4e9cbcaf2c Mon Sep 17 00:00:00 2001 From: aiXander Date: Mon, 15 Apr 2024 08:36:18 -0700 Subject: [PATCH] remove trainable embeddings from print --- trainer/models.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/trainer/models.py b/trainer/models.py index 5181e2e..31688d8 100644 --- a/trainer/models.py +++ b/trainer/models.py @@ -111,9 +111,9 @@ def download_weights(url, dest): def print_trainable_parameters(model, name = ''): trainable_params = 0 all_param = 0 - for _, param in model.named_parameters(): + for name, param in model.named_parameters(): all_param += param.numel() - if param.requires_grad: + if param.requires_grad and "token_embedding" not in name: trainable_params += param.numel() line_delimiter = "#" * 80 print(line_delimiter)