Thank you for the response! I do have it working as a custom Python module currently and have some initial results that show potential. From what we have seen, we are wanting to explore further. My advisor is wanting me to modify the source C++ code so that whatever we would present in a paper is exactly the behavior one would expect if our method was implemented into PyTorch and another person uses it.
I’m curious if I were to rewrite the normal PyTorch batch norm layer using the PyTorch CPP extension 9, should I be able to exactly match the performance of PyTorch’s version? If that’s the case, it seems like it’d be easier to implement it that way, but we would also be able to get close enough to what my advisor is wanting.