docstr: pre_trained: Bool Variable to check whether pre trained model exists to avoid training again
docstr_tok: ["pre_trained",":","Bool","Variable","to","check","whether","pre","trained","model","exists","to","avoid","training","again"]
code: def train(self, pre_trained=False):
      """
      pre_trained: Bool Variable to check whether pre trained model exists to avoid training again
      """
      if pre_trained:
          self.cf_vae.load_state_dict(torch.load(self.save_path))
          self.cf_vae.eval()
          return
  
      # TODO: Handling such dataset specific constraints in a more general way
      # CF Generation for only low to high income data points
      self.vae_train_dataset = self.vae_train_dataset[
          self.vae_train_dataset[:, -1] == 0, :
      ]
      self.vae_val_dataset = self.vae_val_dataset[
          self.vae_val_dataset[:, -1] == 0, :
      ]
  
      # Removing the outcome variable from the datasets
      self.vae_train_feat = self.vae_train_dataset[:, :-1]
      self.vae_val_feat = self.vae_val_dataset[:, :-1]
  
      for epoch in range(self.epochs):
          batch_num = 0
          train_loss = 0.0
          train_size = 0
  
          train_dataset = torch.tensor(self.vae_train_feat).float()
          train_dataset = torch.utils.data.DataLoader(
              train_dataset, batch_size=self.batch_size, shuffle=True
          )
          for train in enumerate(train_dataset):
              self.cf_vae_optimizer.zero_grad()
  
              train_x = train[1]
              train_y = 1.0 - torch.argmax(
                  self.pred_model(train_x), dim=1
              )
              train_size += train_x.shape[0]
  
              out = self.cf_vae(train_x, train_y)
              loss = self.compute_loss(out, train_x, train_y)
              loss.backward()
              train_loss += loss.item()
              self.cf_vae_optimizer.step()
  
              batch_num += 1
  
          ret = loss / batch_num
          print("Train Avg Loss: ", ret, train_size)
  
          # Save the model after training every 10 epochs and at the last epoch
          if (
              epoch != 0 and (epoch % 10) == 0
          ) or epoch == self.epochs - 1:
              torch.save(self.cf_vae.state_dict(), self.save_path)
  
file_path: DiCE/dice_ml/explainer_interfaces/feasible_base_vae.py
imports: import numpy
  import torch
  import torch.utils.data
  from torch.nn import functional
  from dice_ml import diverse_counterfactuals
  from dice_ml.explainer_interfaces.explainer_base import ExplainerBase
  from dice_ml.utils.helpers import get_base_gen_cf_initialization
label_desc: Use image node_id 3 for calling the FeasibleBaseVAE obj's underlying member method code with example usage: obj.train(pre_trained) without return types