Fixing MixUp and CutMix error

[meinkappa]

Hello, all. I am new to programming, and I found an issue with fastai that I might be able to help.

There is an issue #3486 on fastai regarding MixUp and CutMix:

CutMix or MixUp doesnt seem to work with training after using lr_finder

I think this error is happening because MixHandler callback runs before_train twice, which runs
self.old_lf,self.learn.loss_func = self.learn.loss_func,self.lf
twice.
It is running twice because both CutMix and MixUp have it. When it tries to run the second time, getattr(self.learn.loss_func, 'y_int', False) returns False, and the RuntimeError happens.

It can be fixed by changing MixHandler like this:

class MixHandler(Callback):
    "A handler class for implementing `MixUp` style scheduling"
    run_valid = False
    def __init__(self, alpha=0.5):
        self.swapped = False
        self.distrib = Beta(tensor(alpha), tensor(alpha))

    def before_train(self):
        self.stack_y = getattr(self.learn.loss_func, 'y_int', False)
        if self.stack_y and not self.swapped:
            self.old_lf,self.learn.loss_func = self.learn.loss_func,self.lf
            self.swapped = True

    def after_train(self):
        if self.stack_y and self.swapped: 
            self.learn.loss_func = self.old_lf
            self.swapped = False

    def after_cancel_train(self):
        self.after_train()

    def lf(self, pred, *yb):
        if not self.training: return self.old_lf(pred, *yb)
        with NoneReduce(self.old_lf) as lf:
            loss = torch.lerp(lf(pred,*self.yb1), lf(pred,*yb), self.lam)
        return reduce_loss(loss, getattr(self.old_lf, 'reduction', 'mean'))
    
    lf.y_int = True if getattr(learn.loss_func, 'y_int', False) or getattr(lr, 'y_int', False) else False

Can anyone take a look at this and see whether it works on others as well?
And if it does, is it an acceptable style for using Python on fastai standard?
I am new to Python and fastai, so anything does not sound right, please let me know.
Thanks.