Sure! No difference and the shape is as expected:
CUDA_VISIBLE_DEVICES=0,1 python tmp.py
torch.Size([64, 32]) # [batch_size, batch_size/n_gpu]

using:
import torch
import torch.nn as nn
import torchvision.models as models

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.layer = models.resnext50_32x4d(pretrained = False)
        self.layer.fc = nn.Linear(2048,100)
        self.query = nn.Linear(100,100, bias = False)
        self.key = nn.Linear(100,100, bias = False)
        self.value = nn.Linear(100,100, bias = False)
        self.nl = nn.LogSoftmax(dim = 1)

    def forward(self, X):
        x = self.layer(X)
        q = self.query(x)
        k = self.key(x)
        v = self.value(x)
        kqdots = torch.matmul(q, torch.transpose(k, 0, 1))
        return kqdots

model = Model()
model = nn.DataParallel(model).cuda()
x = torch.randn(64, 3, 256 ,256).cuda()
out = model(x)
print(out.shape)