fix: fixed model shapes + redit training loop

This commit is contained in:
Robin Meersman 2025-11-27 14:11:53 +01:00
parent ed44d5b283
commit eb4a014aa1
3 changed files with 68 additions and 56 deletions

View file

@ -12,19 +12,13 @@ from train import train
def create_model(trial: tr.Trial, vocab_size: int = 256):
num_layers = trial.suggest_int("num_layers", 1, 6)
hidden_dim = trial.suggest_int("hidden_dim", 64, 512, log=True)
kernel_size = trial.suggest_int("kernel_size", 2, 7)
dropout_prob = trial.suggest_float("dropout_prob", 0.1, 0.5)
use_batchnorm = trial.suggest_categorical("use_batchnorm", [True, False])
embedding_dim = trial.suggest_int("embedding_dim", 64, 512, log=True)
return CNNPredictor(
vocab_size=vocab_size,
num_layers=num_layers,
hidden_dim=hidden_dim,
kernel_size=kernel_size,
dropout_prob=dropout_prob,
use_batchnorm=use_batchnorm
embed_dim=embedding_dim,
)