How can I more efficiently multiply every element in a batch of tensors with every other batch element, except for itself?
So, I have this code that multiplies every element in a batch of tensors with every other element, except for itself. The code works, but it becomes painfully slow with larger batch sizes (Ideally I want to be able to use it with batch sizes of up to 1000 or more, but even a couple hundred is okay). It basically freezes when using the PyTorch autograd system and large batch sizes (like 50 or greater).
I need help making the code faster and more efficient, while still getting the same output. Any help would be appreciated!
import torch

tensor = torch.randn(50, 512, 512)

batch_size = tensor.size(0)
list1 = []
for i in range(batch_size):
    list2 = []
    for j in range(batch_size):
        if j != i:
            x_out = (tensor[i] * tensor[j]).sum()
            list2.append(x_out )
    list1.append(sum(list2))

out = sum(list1)

I thought that torch.prod might be able to be used, but it doesn't seem to result in the same output as the code above. NumPy answers are acceptable as long as they can be recreated in PyTorch.

You could do the following:
import torch
tensor = torch.randn(50, 512, 512)

batch_size = tensor.size(0)
tensor = tensor.reshape(batch_size, -1)
prod = torch.matmul(tensor, tensor.transpose(0,1))
out = torch.sum(prod) - torch.trace(prod)

Here, you first flatten each element. Then, you multiply the matrix where each row is an element with its own transpose, which gives a batch_size x batch_size matrix, where the ijth element equals the product of tensor[i] with tensor[j]. So, summing up over the values in this matrix and subtracting its trace (i.e., sum of diagonal elements) gives the desired result.
I tried both methods with a batch_size of 1000, and the time taken dropped from 61.43s to 0.59s.
