class CustAct(torch.nn.Module):
    def __init__(self, alphas=1):
        super(CustAct, self).__init__()
        self.alphas = torch.nn.Parameter(torch.full((1,), alphas))
        self.register_buffer("scale", torch.tensor([alphas]) #if you want to separate non-trainable constant, multiply in forward()
        self.scale = alphas #works too

     def forward(self, x):
        val = customAutogradFunction(x, self.alphas) #if you must provide custom gradients
        val = x * self.alphas #no need for autograd.Function
        val = x.clamp(min=0) #your MyReLU as is also needs no autograd.Function
        return val




 silvester:

Is there a way that I can remove cuda() so it can automatically differentiate between cpu and cuda? Even when I use CustAct.cuda(), still I need to explicitly map self.alphas to GPU.


that’s because of the way you used multiplication, your self.alphas became Tensor not nn.Parameter and was not moved. Another idiom is tensorA2=tensorA.to(tensorB), this changes device and dtype, but it is not usually needed.