class combine_act(nn.Module):
     def __init__(self,in_features,comb='convex'):
          super(combine_act, self).__init__()
    
    self.in_features = in_features
    self.comb = comb                     #type of combination [affine or convex]
    
    Lin = lambda x: x

    self.activations = {Lin, nn.ReLU(), nn.Tanh()}           #list of activations to use 

    self.alpha = nn.Parameter(torch.randn(len(self.activations)))
    self.alpha.requiresGrad = True
    
    


    def my_softmax(self, x):
      means = torch.mean(x)
      x_exp = torch.exp(x-means)
      x_exp_sum = torch.sum(x_exp)
      return x_exp/x_exp_sum

    def unitnorm(self,x):
      eps = 1e-7
      return x / x.sum().clamp(min=eps)

   def get_alpha(self,a):
      if self.comb=='affine':
         return self.unitnorm(a).cuda()
      else:
         return self.my_softmax(a).cuda()

   def forward(self, x):
      out=x-x 
      weight = self.get_alpha(self.alpha)
      for i, activation in enumerate(self.activations):
         out += weight[i]*activation(x)
         return out


@ptrblck I have a working prototype now. I have to do some extensive tests. (see code above)
But there still remains an issue with the  update of weights.
In forward function I am computing weights from trainable parameter alpha. However, What I want is to update alpha directly i.e., do gradient update on alpha via autograd and then rescale values according to convex/affine constraints so that I have the correct model parameter.
Any suggestions?
Thanks