class_docstr: This class tests the crypten.nn package.
class_docstr_tok: ["This","class","tests","the","crypten.nn","package","."]
docstr: Tests training of simple model in crypten.nn.
code: def test_training(self):
      """
      Tests training of simple model in crypten.nn.
      """
  
      # create MLP with one hidden layer:
      learning_rate = 0.1
      batch_size, num_inputs, num_intermediate, num_outputs = (
          8,
          10,
          5,
          1,
      )
      model = crypten.nn.Sequential(
          crypten.nn.Linear(num_inputs, num_intermediate),
          crypten.nn.ReLU(),
          crypten.nn.Linear(num_intermediate, num_outputs),
      )
      model.train()
      model.encrypt()
      loss = crypten.nn.MSELoss()
  
      # perform training iterations:
      for _ in range(10):
          for compute_gradients in [True, False]:
              # get training sample:
              input = get_random_test_tensor(
                  size=(batch_size, num_inputs), is_float=True
              )
              target = input.mean(dim=1, keepdim=True)
  
              # encrypt training sample:
              input = crypten.cryptensor(input)
              target = crypten.cryptensor(target)
              if compute_gradients:
                  input.requires_grad = True
                  target.requires_grad = True
  
              # perform forward pass:
              output = model(input)
              loss_value = loss(output, target)
  
              # set gradients to "zero" (setting to None is more efficient):
              model.zero_grad()
              for param in model.parameters():
                  self.assertIsNone(
                      param.grad, "zero_grad did not reset gradients"
                  )
  
              # perform backward pass:
              loss_value.backward()
  
              # perform parameter update:
              reference = {}
              reference = self._compute_reference_parameters(
                  "", reference, model, learning_rate
              )
              model.update_parameters(learning_rate)
              self._check_reference_parameters("", reference, model)
  
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 19 for calling the TestNN obj's underlying member method code with example usage: obj.test_training() without return types