How to set training weight on softmax during training

Question

I'm reproducing Auto-DeepLab with PyTorch and I got a problem, that is, I can't set the architecture weight(both cell and layer) on softmax. It either leads to twice backward or weights get no upgrade with gradients but only softmax.

class Architect () :
    def __init__(self, model, args):
        self.network_momentum = args.momentum
        self.network_weight_decay = args.weight_decay
        self.model = model
        self.optimizer = 
            torch.optim.Adam(self.model.arch_parameters(),
                lr=args.arch_lr, betas=(0.5, 0.999), 
                weight_decay=args.arch_weight_decay)

    def step (self, input_valid, target_valid) :
        self.model.soft_parameters()
        self.optimizer.zero_grad ()
        self._backward_step(input_valid, target_valid)
        self.optimizer.step()

    def _backward_step (self, input_valid, target_valid) :
        _, loss = self.model._loss (input_valid, target_valid)
        loss.backward ()

under this code, it leads softmax works but the weights are not optimized with gradients.

[[0.3333, 0.3333, 0.3333],
         [0.3333, 0.3333, 0.3333],
         [0.3333, 0.3333, 0.3333],
         [0.3333, 0.3333, 0.3333]]*12], device='cuda:0',
       grad_fn=<SoftmaxBackward>)