code: def __init__(
      self,
      num_steps,
      x_size,
      window_size,
      z_what_size,
      rnn_hidden_size,
      encoder_net=[],
      decoder_net=[],
      predict_net=[],
      embed_net=None,
      bl_predict_net=[],
      non_linearity="ReLU",
      decoder_output_bias=None,
      decoder_output_use_sigmoid=False,
      use_masking=True,
      use_baselines=True,
      baseline_scalar=None,
      scale_prior_mean=3.0,
      scale_prior_sd=0.1,
      pos_prior_mean=0.0,
      pos_prior_sd=1.0,
      likelihood_sd=0.3,
      use_cuda=False,
  ):
      super().__init__()
  
      self.num_steps = num_steps
      self.x_size = x_size
      self.window_size = window_size
      self.z_what_size = z_what_size
      self.rnn_hidden_size = rnn_hidden_size
      self.use_masking = use_masking
      self.use_baselines = use_baselines
      self.baseline_scalar = baseline_scalar
      self.likelihood_sd = likelihood_sd
      self.use_cuda = use_cuda
      prototype = (
          torch.tensor(0.0).cuda() if use_cuda else torch.tensor(0.0)
      )
      self.options = dict(
          dtype=prototype.dtype, device=prototype.device
      )
  
      self.z_pres_size = 1
      self.z_where_size = 3
      # By making these parameters they will be moved to the gpu
      # when necessary. (They are not registered with pyro for
      # optimization.)
      self.z_where_loc_prior = nn.Parameter(
          torch.FloatTensor(
              [scale_prior_mean, pos_prior_mean, pos_prior_mean]
          ),
          requires_grad=False,
      )
      self.z_where_scale_prior = nn.Parameter(
          torch.FloatTensor(
              [scale_prior_sd, pos_prior_sd, pos_prior_sd]
          ),
          requires_grad=False,
      )
  
      # Create nn modules.
      rnn_input_size = (
          x_size**2 if embed_net is None else embed_net[-1]
      )
      rnn_input_size += (
          self.z_where_size + z_what_size + self.z_pres_size
      )
      nl = getattr(nn, non_linearity)
  
      self.rnn = nn.LSTMCell(rnn_input_size, rnn_hidden_size)
      self.encode = Encoder(
          window_size**2, encoder_net, z_what_size, nl
      )
      self.decode = Decoder(
          window_size**2,
          decoder_net,
          z_what_size,
          decoder_output_bias,
          decoder_output_use_sigmoid,
          nl,
      )
      self.predict = Predict(
          rnn_hidden_size,
          predict_net,
          self.z_pres_size,
          self.z_where_size,
          nl,
      )
      self.embed = (
          Identity()
          if embed_net is None
          else MLP(x_size**2, embed_net, nl, True)
      )
  
      self.bl_rnn = nn.LSTMCell(rnn_input_size, rnn_hidden_size)
      self.bl_predict = MLP(rnn_hidden_size, bl_predict_net + [1], nl)
      self.bl_embed = (
          Identity()
          if embed_net is None
          else MLP(x_size**2, embed_net, nl, True)
      )
  
      # Create parameters.
      self.h_init = nn.Parameter(torch.zeros(1, rnn_hidden_size))
      self.c_init = nn.Parameter(torch.zeros(1, rnn_hidden_size))
      self.bl_h_init = nn.Parameter(torch.zeros(1, rnn_hidden_size))
      self.bl_c_init = nn.Parameter(torch.zeros(1, rnn_hidden_size))
      self.z_where_init = nn.Parameter(
          torch.zeros(1, self.z_where_size)
      )
      self.z_what_init = nn.Parameter(torch.zeros(1, self.z_what_size))
  
      if use_cuda:
          self.cuda()
  
imports: from collections import namedtuple
  import torch
  import torch.nn
  import torch.nn.functional
  from modules import MLP, Decoder, Encoder, Identity, Predict
  import pyro
  import pyro.distributions
label_desc: Use image node_id 1 to create a new AIR object from inherited base classes: nn with example: obj = AIR(num_steps, x_size, window_size, z_what_size, rnn_hidden_size, encoder_net, decoder_net, predict_net, embed_net, bl_predict_net, non_linearity, decoder_output_bias, decoder_output_use_sigmoid, use_masking, use_baselines, baseline_scalar, scale_prior_mean, scale_prior_sd, pos_prior_mean, pos_prior_sd, likelihood_sd, use_cuda)