fix keyerror?

This commit is contained in:
Cyberes 2023-10-20 14:00:24 -06:00
parent 0abd4b94fb
commit e838f591aa
1 changed files with 8 additions and 6 deletions

View File

@ -5,9 +5,11 @@ def estimate_model_size(config: dict):
:param config:
:return:
"""
vocab_size = config['vocab_size']
hidden_size = config['hidden_size']
num_hidden_layers = config['num_hidden_layers']
intermediate_size = config['intermediate_size']
total_params = (vocab_size * hidden_size) + (num_hidden_layers * ((hidden_size * intermediate_size * 4) + (hidden_size * hidden_size * 3)))
return int(total_params / 1e9)
vocab_size = config.get('vocab_size')
hidden_size = config.get('hidden_size')
num_hidden_layers = config.get('num_hidden_layers')
intermediate_size = config.get('intermediate_size')
if vocab_size and hidden_size and num_hidden_layers and intermediate_size:
total_params = (vocab_size * hidden_size) + (num_hidden_layers * ((hidden_size * intermediate_size * 4) + (hidden_size * hidden_size * 3)))
return int(total_params / 1e9)
return 0