Is there a way to select different loss functions for different targets? about fastai_dev HOT 2 CLOSED

class TargetLoss():
    def __init__(self, loss_func, target_index, name=None):
        store_attr(self, "loss_func,target_index,name")
    def __call__(self, x, *targets, **kwargs):
        return self.loss_func(x, targets[self.target_index], **kwargs)
    @property
    def __name__(self):
        return self.name
    
class SwitchLoss(Callback):
    def __init__(self, train_loss, valid_loss):
        store_attr(self, "train_loss,valid_loss")
    def begin_train(self):
        self.learn.loss_func = self.train_loss
    def begin_validate(self):
        self.learn.loss_func = self.valid_loss    

train_loss = TargetLoss(MSELossFlat(), 1)
valid_loss = TargetLoss(CrossEntropyLossFlat(), 0, name="valid_loss")
valid_accuracy = TargetLoss(accuracy, 0, name="valid_accuracy")

opt_func = partial(Adam, wd=0.1, eps=1e-7)
cb_funcs = [partial(RNNTrainer, alpha=2, beta=1), partial(SwitchLoss, train_loss, valid_loss)]

learn = Learner(dbch, model, loss_func=train_loss, opt_func=opt_func, cb_funcs=cb_funcs, metrics=[valid_accuracy])