Commit 7c9330b4 authored by narugo1992's avatar narugo1992
Browse files

dev(narugo): fix a bug on loss function

parent 9e7406aa

zoo/monochrome/train_.py

+1 −1
Original line number Diff line number Diff line
@@ -118,7 +118,7 @@ def train(dataset_dir: str, session_name: Optional[str] = None, from_ckpt: Optio 
        loss_weight = torch.as_tensor([torch.e, 1.0]) ** -preference

    else:

        loss_weight = torch.as_tensor([1.0, torch.e]) ** preference

    loss_fn = nn.CrossEntropyLoss(weight=loss_weight)

    loss_fn = nn.CrossEntropyLoss(weight=loss_weight).to(device)

    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)

    scheduler = lr_scheduler.OneCycleLR(

        optimizer, max_lr=learning_rate,