Trying to implement a local-gram loss

I am trying to reproduce a paper and I have a question when trying to implement the local gram loss. Global Gram loss is implemented in the online tutorial of style (transferhttps://pytorch.org/tutorials/advanced/neural_style_tutorial.html), and a local gram is implemented to every small sliding window of the image. To provide a detailed description: an nn.functional.conv2d can be considered as a weighted summation applied to every sliding window with kernel weights, and a local gram can be considered as a gram matrix loss applied to every sliding window with kernel size, and losses are then averaged among all channels and all pixels.
Untitled

It is definitely possible to implement by creating a nn.functional called ‘local-gram’ and derive forward and backward for it. But is it possible to assemble one with existing nn.functionals? I think it should be something similar to pandas.DataFrame.apply but didn’t find it in the documentations? Thank you.