class_docstr: This class tests the crypten.nn package.
class_docstr_tok: ["This","class","tests","the","crypten.nn","package","."]
code: def test_distances(self):
      distances = ["CosineSimilarity"]
  
      for distance in distances:
          x1 = get_random_test_tensor(is_float=True)
          x2 = get_random_test_tensor(is_float=True)
          x1.requires_grad = True
          x2.requires_grad = True
  
          x1_enc = crypten.cryptensor(x1, requires_grad=True)
          x2_enc = crypten.cryptensor(x2, requires_grad=True)
  
          dist_fn = getattr(torch.nn, distance)()
          enc_dist_fn = getattr(crypten.nn, distance)()
  
          # Forward Pass
          dist = dist_fn(x1, x2)
          dist_enc = enc_dist_fn(x1_enc, x2_enc)
  
          self._check(dist_enc, dist, f"{distance} failed in forward")
  
          # Backward Pass
          grad_output = get_random_test_tensor(
              size=dist.size(), is_float=True
          ).abs()
          grad_output_enc = crypten.cryptensor(grad_output)
  
          dist.backward(grad_output)
          dist_enc.backward(grad_output_enc)
  
          self._check(
              x1_enc.grad, x1.grad, f"{distance} 1st arg grad failed"
          )
          self._check(
              x2_enc.grad, x2.grad, f"{distance} 2nd arg grad failed"
          )
  
imports: import copy
  import itertools
  import logging
  import unittest
  import crypten
  import crypten.communicator
  import torch
  import torch.nn.functional
  from crypten.common.rng import generate_random_ring_element
  from crypten.common.tensor_types import is_float_tensor
  from crypten.config import cfg
  from crypten.encoder import FixedPointEncoder
  from test.multiprocess_test_case import get_random_linear, get_random_test_tensor, MultiProcessTestCase, onehot
label_desc: Use image node_id 17 for calling the TestNN obj's underlying member method code with example usage: obj.test_distances() without return types