Create the labels directly on GPU. When you use .cuda(), you are first creating the tensor on CPU and then moving the tensor to GPU which is slow. So instead use this
labels = torch.tensor(range(q.shape[0]), dtype=torch.int64, device=q.device)

(There is nothing special about using torch.tensor)