albanD:

autograd


Thank you for in-time reply.
I have test your suggestion in Python env by creating a function first as below.
#submodules.py
import numpy as np
import torch
from torch.autograd import Function
class CheckBP(Function):

    def __init__(self, label, show):
        super(CheckBP, self).__init__()
        self.label = label
        self.show = show

    def forward(self, input):
        # print(self.label + ': forward passed.')
        # print(input)
        return input.clone()

    def backward(self, grad_output):
        grad_mean = grad_output.abs().mean()
        if self.show == 1:
            print('grad_' + self.label + ': ' + str(grad_mean))
        return grad_output

Then,  a class is used to wrap the function as below:
# modules
import numpy as np
import torch
import torch.nn as nn
from torch.autograd import Variable
import submodules as F

class CheckBP(nn.Module):

    def __init__(self, label='a', show=1):
        super(CheckBP, self).__init__()
        self.check_bp = F.CheckBP(label, show)

    def forward(self, input):
        return self.check_bp(input)

Testing the CheckBP in Python is ok.
However, in c++(libtorch),
class checkBP :public torch::autograd::Function
{
...(something)
}

I got this error " error: invalid use of incomplete type ‘struct torch::autograd::Function
struct checkBP : public torch::autograd::Function "
I checked “torch::autograd::Function” in   /torch/csrc/autograd/edge.h . There is forward() or backward() to inherit.
So, how to define the checkBP module in c++ like in python with self-defined backward()?
Thank for your reply again!