//  [python]
def train(self, dataset_train, dataset_val, model_name, model_folder):
        self.balancing_representation = self.build_balancing_representation()
        self.treatment_prob_predictions = self.build_treatment_assignments_one_hot(self.balancing_representation)
        self.predictions = self.build_outcomes(self.balancing_representation)

        self.loss_treatments = self.compute_loss_treatments_one_hot(target_treatments=self.current_treatments,
                                                                    treatment_predictions=self.treatment_prob_predictions,
                                                                    active_entries=self.active_entries)
        self.loss_outcomes = self.compute_loss_predictions(self.outputs, self.predictions, self.active_entries)
        self.loss = self.loss_outcomes + self.loss_treatments
        optimizer = self.get_optimizer()

        # Setup tensorflow
        tf_device = 'gpu'
        if tf_device == "cpu":
            tf_config = tf.ConfigProto(log_device_placement=False, device_count={'GPU': 0})
        else:
            tf_config = tf.ConfigProto(log_device_placement=False, device_count={'GPU': 1})
            tf_config.gpu_options.allow_growth = True

        self.sess = tf.Session(config=tf_config)
        self.sess.run(tf.global_variables_initializer())
        self.sess.run(tf.local_variables_initializer())

        for epoch in range(self.num_epochs):
            p = float(epoch) / float(self.num_epochs)
            alpha_current = 2. / (1. + np.exp(-10. * p)) - 1

            iteration = 0
            for (batch_current_covariates, batch_previous_treatments, batch_current_treatments, batch_init_state,
                 batch_outputs, batch_active_entries) in self.gen_epoch(dataset_train, batch_size=self.batch_size):
                feed_dict = self.build_feed_dictionary(batch_current_covariates, batch_previous_treatments,
                                                       batch_current_treatments, batch_init_state, batch_outputs,
                                                       batch_active_entries,
                                                       alpha_current)

                _, training_loss, training_loss_outcomes, training_loss_treatments = self.sess.run(
                    [optimizer, self.loss, self.loss_outcomes, self.loss_treatments],
                    feed_dict=feed_dict)

                iteration += 1

            logging.info(
                "Epoch {} out of {} | total loss = {} | outcome loss = {} | "
                "treatment loss = {} | current alpha = {} ".format(epoch + 1, self.num_epochs, training_loss,
                                                                   training_loss_outcomes,
                                                                   training_loss_treatments, alpha_current))

        # Validation loss
        validation_loss, validation_loss_outcomes, \
        validation_loss_treatments = self.compute_validation_loss(dataset_val)

        validation_mse, _ = self.evaluate_predictions(dataset_val)

        logging.info(
            "Epoch {} Summary| Validation total loss = {} | Validation outcome loss = {} | Validation treatment loss {} | Validation mse = {}".format(
                epoch, validation_loss, validation_loss_outcomes, validation_loss_treatments, validation_mse))

        checkpoint_name = model_name + "_final"
        self.save_network(self.sess, model_folder, checkpoint_name)