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 (transfer, 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.

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.