When `self.update_weights_func(effective_batch_loss)`, should we do like this? ` self.optimizer.zero_grad() self.update_weights_func(effective_batch_loss) self.optimizer.step() `