diff --git a/src/ptbench/configs/models/densenet_pretrained.py b/src/ptbench/configs/models/densenet_pretrained.py
index cb729083ab89fa96e13232ad7223c1659d3cdc69..b018a52203061b847cdae9f09b5edfa713930302 100644
--- a/src/ptbench/configs/models/densenet_pretrained.py
+++ b/src/ptbench/configs/models/densenet_pretrained.py
@@ -4,19 +4,22 @@
 
 """DenseNet."""
 
+from torch import empty
 from torch.nn import BCEWithLogitsLoss
-from torch.optim import Adam
 
-from ...models.densenet import build_densenet
+from ...models.densenet import Densenet
 
 # config
-lr = 0.01
-
-# model
-model = build_densenet(pretrained=True)
+optimizer_configs = {"lr": 0.01}
 
 # optimizer
-optimizer = Adam(model.parameters(), lr=lr)
+optimizer = "Adam"
 
 # criterion
-criterion = BCEWithLogitsLoss()
+criterion = BCEWithLogitsLoss(pos_weight=empty(1))
+criterion_valid = BCEWithLogitsLoss(pos_weight=empty(1))
+
+# model
+model = Densenet(
+    criterion, criterion_valid, optimizer, optimizer_configs, pretrained=True
+)