(1) If, i write “output = torch.square(input)” in customize activation forward() in below code, i think it doesn’t require backward() implementation as it used torch function
(2) But if i write “output = input * input”, it requires the implementation of backward().  Is below code is fine for customize function “output = input * input”?
Let me know if i am missing something.
class my_act(Function):

# Note that both forward and backward are @staticmethods
@staticmethod
 def forward(ctx, input, weight, bias=None):
    ctx.save_for_backward(input)
    output = input * input
    return output

@staticmethod
def backward(ctx, grad_output):
   
    input = ctx.saved_tensors
    grad_input = grad_output * ( 2 * x) # DERIVATIVE OF X*X MULTIPLY WITH GRAD_OUTPUT
  
    return grad_input