Add a custom criterion and optimizer support for utils/segmentation.py
Closed this issue · 0 comments
ishan121028 commented
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>)