Hi Composite!
As an aside, your code is unhelpfully obtuse.



 notprime:

    def function(self, X, Y):
        ...
        Kc = torch.matmul(torch.matmul(H, K), H)
        Lc = torch.matmul(torch.matmul(H, L), H)

        testStat = torch.sum(torch.transpose(Kc, 0, 1) * Lc) / n
        ...



You loop over pairs of columns from data and then call self.function()
on that pair.
But up until the line in which you calculate testStat, Kc depends only
on the first column in the pair and Lc only on the second.  So outside
of your double loop you can precompute everything that goes into Kc
and Lc.
First do this precomputation to clear away all of the confusion.
Then, to “vectorize” the computation of testStat and placing in into the
matrix stat, take a look at torch.einsum().
Good luck!
K. Frank