    def save(self, savedir, timestep):
        """Saves all information required to recreate this model in two files in savedir
        (or self.model_dir if savedir is None), one containing the model structuure and the other
        containing all variables in the network.

        savedir (str): (Optional) Path to which files will be saved. If not provided, self.model_dir
            (the directory provided at initialization) will be used.
        """
        if not self.finalized:
            raise RuntimeError()
        model_dir = self.model_dir if savedir is None else savedir

        # Write structure to file
        with open(os.path.join(model_dir, '{}_{}.nns'.format(self.name, timestep)), "w+") as f:
            for layer in self.layers[:-1]:
                f.write("%s\n" % repr(layer))
            last_layer_copy = self.layers[-1].copy()
            last_layer_copy.set_activation(self.end_act_name)
            last_layer_copy.set_output_dim(last_layer_copy.get_output_dim() // 2)
            f.write("%s\n" % repr(last_layer_copy))

        # Save network parameters (including scalers) in a .mat file
        var_vals = {}
        for i, var_val in enumerate(self.sess.run(self.nonoptvars + self.optvars)):
            var_vals[str(i)] = var_val
        savemat(os.path.join(model_dir, '{}_{}.mat'.format(self.name, timestep)), var_vals)
Saves all information required to recreate this model in two files in savedir
(or self.model_dir if savedir is None), one containing the model structuure and the other
containing all variables in the network.

savedir (str): (Optional) Path to which files will be saved. If not provided, self.model_dir
(the directory provided at initialization) will be used.
mbpo/models/bnn.py