Out of Memory

[PHChenGit]

Hi,

I'm trying to reproduce your work with my GPU RTX 4090.
But I got Error about GPU memory out of memory when computing bmm()
Do you have any suggest to fix or avoid this error?

Thanks for your help.

def compute_similarity(features_a, features_b):
    b, c, h, w = features_a.shape
    features_a = features_a.transpose(2, 3).contiguous().view(b, c, h*w)
    features_b = features_b.view(b, c, h*w).transpose(1, 2)
    features_mul = torch.bmm(features_b, features_a)
    correlation_tensor = features_mul.view(b, h, w, h*w).transpose(2, 3).transpose(1, 2)
    correlation_tensor = F.normalize(F.relu(correlation_tensor), p=2, dim=1)
    return correlation_tensor