import torch

def guide_output(input, guide):
    x = input.squeeze(0)
    y = guide.squeeze(0)
    ch = x.shape[0]
    x = x.reshape(ch,-1)
    y = y.reshape(ch,-1)
    print(x.T.shape, y.shape)
    A = torch.mm(x.T, y) # compute the matrix of dot-products with guide features
    x.reshape(ch,-1)[:] = y[:,A.argmax(1)] # select ones that match best

input = torch.randn(3,512,512).unsqueeze(0)
guide = torch.randn(3,512,512).unsqueeze(0)
output = guide_output(input, guide)

The printed shape is torch.Size([262144, 3]) torch.Size([3, 262144]). Thus the size of the output of your torch.mm is a tensor of shape (262144,262144) which requires
274877906944 bytes i.e. 274.877906944 Gigabytes. And since the system can’t provide it, we see the error.