tdpeuter/2025ML-project-neural_compression
tdpeuter/2025ML-project-neural_compression
Archived
1
Fork 0
Code Issues Pull requests Releases 1 Activity Actions

fix: Properly pass device

Browse source
This commit is contained in:
Tibo De Peuter 2025-12-09 14:50:25 +01:00
parent 28ae8191ad
commit 8311eabd4d
Signed by: tdpeuter
GPG key ID: 38297DE43F75FFE2
3 changed files with 2 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

2
src/trainers/FullTrainer.py
View file

@ -20,7 +20,7 @@ class FullTrainer(Trainer):
raise ValueError("Model must be provided: run optuna optimizations first") raise ValueError("Model must be provided: run optuna optimizations first")
model.to(device) model.to(device)
train_loss, val_loss = train(model, train_loader, validation_loader, model.loss_function, n_epochs) train_loss, val_loss = train(model, train_loader, validation_loader, model.loss_function, n_epochs, device=device)
print_losses(train_loss, val_loss) print_losses(train_loss, val_loss)
return model return model

2
src/trainers/OptunaTrainer.py
View file

@ -39,7 +39,7 @@ def objective_function(
device: str device: str
): ):
model = create_model(trial, model).to(device) model = create_model(trial, model).to(device)
_, validation_loss = train(model, training_loader, validation_loader, model.loss_function) _, validation_loss = train(model, training_loader, validation_loader, model.loss_function, device=device)
return min(validation_loss) return min(validation_loss)

1
src/trainers/train.py
View file

@ -15,7 +15,6 @@ def train(
weight_decay: float = 1e-8, weight_decay: float = 1e-8,
device="cuda" device="cuda"
) -> tuple[list[float], list[float]]: ) -> tuple[list[float], list[float]]:
model.to(device) model.to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay) optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)