You can use this snippet. Cosine similarity is the same as the scalar product of the normalized inputs and you can get the pw scalar product through matrix multiplication.
Cosine distance in turn is just 1-cosine_similarity.
def pw_cosine_distance(input_a, input_b):
   normalized_input_a = torch.nn.functional.normalize(input_a)  
   normalized_input_b = torch.nn.functional.normalize(input_b)
   res = torch.mm(normalized_input_a, normalized_input_b.T)
   res *= -1 # 1-res without copy
   res += 1
   return res