Dear ptrblck,
I have another problem with this function. Do you think the activation function I built could be auto-grad?
class MultiWavelet(nn.Module):

    def __init__(self,out_features):
        super(MultiWavelet, self).__init__()
        self.out_features = out_features
        self.scale = nn.Parameter(torch.randn(1,self.out_features),requires_grad = True)
        self.shift = nn.Parameter(torch.randn(1,self.out_features),requires_grad = True)
        self.a = nn.Parameter(torch.randn(1,self.out_features),requires_grad = True)
        self.b = nn.Parameter(torch.randn(1,self.out_features),requires_grad = True)
        
    def Wavelet1(self,c):
        self.c = c
        out = self.c.detach()
        out[(out>=0) & (out <1)] = -2 * torch.pow(out[(out>=0) & (out <1)],3) + 3 * torch.pow(out[(out>=0) & (out <1)],2)
        out[(out>=1) & (out <=2)] = torch.pow((2-out[(out>=1) & (out <=2)]),2) * (2*out[(out>=1) & (out <=2)] - 1)
        out[(out <0) | (out>2)] = 0
        return out
        
    
    def Wavelet2(self,c):
        self.c = c
        out = self.c.detach()
        out[(out>=0) & (out <1)] = -torch.pow(out[(out>=0) & (out <1)],2) * (3 * out[(out>=0) & (out <1)] - 3)
        out[(out>=1) & (out <=2)] = -torch.pow((2 - out[(out>=1) & (out <=2)]),2) * (3* out[(out>=1) & (out <=2)] - 3)
        out[(out<0) | (out>2)] = 0
        return out
    
    def forward(self, x):
        x_out = (x-self.shift)/self.scale
        h1 = self.Wavelet1(x_out)
        h2 = self.Wavelet2(x_out)
        h = self.a * h1 + self.b * h2
        
        return h

I need to compute these parameters, such as scale, shift, a, b, etc.
Look forward to hearing from you.
Best regards.