I think I’ve got working with this:
flattened = input.view(input.size(0), input.size(1), -1)
grams = torch.matmul(flattened, torch.transpose(flattened, 1, 2))
list2 = []
for i in range(input.size(0)):
    list1 = []
    for j in range(input.size(0)):
        if j!=i:
            similarity = sum(torch.cosine_similarity(grams[i].view(1,-1), grams[j].view(1,-1)))
            list1.append(similarity)
    list2.append(sum(list1))
loss = -sum(list2) / input.size(0)

Or the last bit rewritten in list comprehension:
loss = -sum([ sum([(torch.cosine_similarity(grams[j].view(1,-1), grams[i].view(1,-1))).sum() for i in range(input.size(0)) if i != j]) for j in range(input.size(0))]) / input.size(0)