wncc/UniTrain

Add a custom criterion and optimizer support for utils/segmentation.py

Closed this issue · 0 comments

Adding support for choosing custom loss functions and optimizers would allow users more flexibility while working out with their experiments.

You have to edit the train_model() function for this:

A solution for this issue would look somethinig like this:

train_model(model=model, train_dataloader=train_dataloader, test_dataloader=test_dataloader, ...... optimizer=torch.optim.<Optimizer>(), criterion=torch.nn.<Criterion>)