def get_cosine_distance(x: Tensor, y: Tensor):
  """
  :param x: [B1, D]
  :param y: [B2, D]
  :return: [B1, B2]
  """
  x_norm = torch.norm(x, dim=1, keepdim=True)
  y_norm = torch.norm(y, dim=1, keepdim=True)
  return torch.mm(x, y.transpose(0, 1)) / x_norm / y_norm.transpose(0, 1)

get_cosine_distance(objective_vectors[i].reshape..., feat)

I think replacing the torch.norm part in your code with get_cosine_distance would resolve your issue.
Is it what you want?