Hi,
I think you want the following:
class MyModule(nn.Module)
  def __init__(self, flag):
    self.flag = flag
  def forward(self, x):
    return MyFunc.apply(x, flag)

class MyFunc(Function):
  def forward(ctx, x, flag):
    ctx.flag = flag # Save your flag
    some operations on x
  def backward(ctx, y):
    if ctx.flag == True:
      automatically compute the gradient
    if ctx.flag == False:
      some operations on the gradient of this layer and then continue to backward

    return gradX, None # Here you want None because no grads are needed for flag.