code: def learn(
      bivae,
      train_set,
      n_epochs,
      batch_size,
      learn_rate,
      beta_kl,
      verbose,
      device=torch.device("cpu"),
      dtype=torch.float32,
  ):
      user_params = it.chain(
          bivae.user_encoder.parameters(),
          bivae.user_mu.parameters(),
          bivae.user_std.parameters(),
      )
  
      item_params = it.chain(
          bivae.item_encoder.parameters(),
          bivae.item_mu.parameters(),
          bivae.item_std.parameters(),
      )
  
      if bivae.cap_priors.get("user", False):
          user_params = it.chain(
              user_params, bivae.user_prior_encoder.parameters()
          )
          user_features = train_set.user_feature.features[
              : train_set.num_users
          ]
  
      if bivae.cap_priors.get("item", False):
          item_params = it.chain(
              item_params, bivae.item_prior_encoder.parameters()
          )
          item_features = train_set.item_feature.features[
              : train_set.num_items
          ]
  
      u_optimizer = torch.optim.Adam(params=user_params, lr=learn_rate)
      i_optimizer = torch.optim.Adam(params=item_params, lr=learn_rate)
  
      x = train_set.matrix.copy()
      x.data = np.ones_like(x.data)  # Binarize data
      tx = x.transpose()
  
      progress_bar = trange(1, n_epochs + 1, disable=not verbose)
      for _ in progress_bar:
          # item side
          i_sum_loss = 0.0
          i_count = 0
          for i_ids in train_set.item_iter(batch_size, shuffle=False):
              i_batch = tx[i_ids, :]
              i_batch = i_batch.A
              i_batch = torch.tensor(
                  i_batch, dtype=dtype, device=device
              )
  
              # Reconstructed batch
              beta, i_batch_, i_mu, i_std = bivae(
                  i_batch, user=False, theta=bivae.theta
              )
  
              i_mu_prior = 0.0  # zero mean for standard normal prior if not CAP prior
              if bivae.cap_priors.get("item", False):
                  i_batch_f = item_features[i_ids]
                  i_batch_f = torch.tensor(
                      i_batch_f, dtype=dtype, device=device
                  )
                  i_mu_prior = bivae.encode_item_prior(i_batch_f)
  
              i_loss = bivae.loss(
                  i_batch, i_batch_, i_mu, i_mu_prior, i_std, beta_kl
              )
              i_optimizer.zero_grad()
              i_loss.backward()
              i_optimizer.step()
  
              i_sum_loss += i_loss.data.item()
              i_count += len(i_batch)
  
              beta, _, i_mu, _ = bivae(
                  i_batch, user=False, theta=bivae.theta
              )
  
              bivae.beta.data[i_ids] = beta.data
              bivae.mu_beta.data[i_ids] = i_mu.data
  
          # user side
          u_sum_loss = 0.0
          u_count = 0
          for u_ids in train_set.user_iter(batch_size, shuffle=False):
              u_batch = x[u_ids, :]
              u_batch = u_batch.A
              u_batch = torch.tensor(
                  u_batch, dtype=dtype, device=device
              )
  
              # Reconstructed batch
              theta, u_batch_, u_mu, u_std = bivae(
                  u_batch, user=True, beta=bivae.beta
              )
  
              u_mu_prior = 0.0  # zero mean for standard normal prior if not CAP prior
              if bivae.cap_priors.get("user", False):
                  u_batch_f = user_features[u_ids]
                  u_batch_f = torch.tensor(
                      u_batch_f, dtype=dtype, device=device
                  )
                  u_mu_prior = bivae.encode_user_prior(u_batch_f)
  
              u_loss = bivae.loss(
                  u_batch, u_batch_, u_mu, u_mu_prior, u_std, beta_kl
              )
              u_optimizer.zero_grad()
              u_loss.backward()
              u_optimizer.step()
  
              u_sum_loss += u_loss.data.item()
              u_count += len(u_batch)
  
              theta, _, u_mu, _ = bivae(
                  u_batch, user=True, beta=bivae.beta
              )
              bivae.theta.data[u_ids] = theta.data
              bivae.mu_theta.data[u_ids] = u_mu.data
  
              progress_bar.set_postfix(
                  loss_i=(i_sum_loss / i_count),
                  loss_u=(u_sum_loss / (u_count)),
              )
  
      # infer mu_beta
      for i_ids in train_set.item_iter(batch_size, shuffle=False):
          i_batch = tx[i_ids, :]
          i_batch = i_batch.A
          i_batch = torch.tensor(i_batch, dtype=dtype, device=device)
  
          beta, _, i_mu, _ = bivae(
              i_batch, user=False, theta=bivae.theta
          )
          bivae.mu_beta.data[i_ids] = i_mu.data
  
      # infer mu_theta
      for u_ids in train_set.user_iter(batch_size, shuffle=False):
          u_batch = x[u_ids, :]
          u_batch = u_batch.A
          u_batch = torch.tensor(u_batch, dtype=dtype, device=device)
  
          theta, _, u_mu, _ = bivae(u_batch, user=True, beta=bivae.beta)
          bivae.mu_theta.data[u_ids] = u_mu.data
  
      return bivae
  
file_path: cornac/cornac/models/bivaecf/bivae.py
imports: import itertools
  import numpy
  import torch
  import torch.nn
  from tqdm.auto import trange
label_desc: Use image node_id 1 for calling a global function with example usage: learn(bivae, train_set, n_epochs, batch_size, learn_rate, beta_kl, verbose, device, dtype) and returns: bivae