Thanks for your answer.
Suppose I need to include a learnable parameter in function autograd,
similar to the function below
class CustAct(torch.nn.Module):
    def __init__(self, alphas=1):
        super(CustAct, self).__init__()
        self.alphas = alphas*torch.nn.Parameter(torch.ones(1))

        self.alphas.requiresGrad = True  

     def forward(self, x):
            val = self.alphas.cuda() * x

        return val

How I can include this in  autograd function?
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.