diff --git a/parameters_combinations/param_combinations_chairs_test.txt b/parameters_combinations/param_combinations_chairs_test.txt
deleted file mode 100644
index 3e0b6a78269c14c316f9469e2fb4f0748db5caca..0000000000000000000000000000000000000000
--- a/parameters_combinations/param_combinations_chairs_test.txt
+++ /dev/null
@@ -1,4 +0,0 @@
---batch-size=64 --dataset=rendered_chairs --epochs=4 --latent_spec_cont=10 --lr=1e-4 --experiment-name=VAE_bs_64_conv_64_64_128_128 --gpu-devices 0 1 --nb-filter-conv1=64 --nb-filter-conv2=64 --nb-filter-conv3=128 --nb-filter-conv4=128
---batch-size=64 --dataset=rendered_chairs --epochs=400 --latent_spec_cont=10 --lr=1e-4 --experiment-name=VAE_bs_64_conv_128_128_256_256 --gpu-devices 0 1 --nb-filter-conv1=128 --nb-filter-conv2=128 --nb-filter-conv3=256 --nb-filter-conv4=256
---batch-size=64 --dataset=rendered_chairs --epochs=400 --latent_spec_cont=30 --lr=1e-4 --experiment-name=VAE_bs_64_conv_64_64_128_128_ls_30 --gpu-devices 0 1 --nb-filter-conv1=64 --nb-filter-conv2=64 --nb-filter-conv3=128 --nb-filter-conv4=128
---batch-size=64 --dataset=rendered_chairs --epochs=400 --latent_spec_cont=30 --lr=1e-4 --experiment-name=VAE_bs_64_conv_128_128_256_256_ls_30 --gpu-devices 0 1 --nb-filter-conv1=128 --nb-filter-conv2=128 --nb-filter-conv3=256 --nb-filter-conv4=256
diff --git a/reconstruction_im/charis_VAE_bs_256.png b/reconstruction_im/charis_VAE_bs_256.png
index bff1fcde28087cc51585bc28b1275d9363135cb8..26d33274ee8245777e8126d6c61e64b4e1bc08cc 100644
Binary files a/reconstruction_im/charis_VAE_bs_256.png and b/reconstruction_im/charis_VAE_bs_256.png differ
diff --git a/reconstruction_im/charis_VAE_bs_64.png b/reconstruction_im/charis_VAE_bs_64.png
index 98cfca4520d226bc7aaf53ce7dcd1b8deaed3200..cf1e1c2900cfc10443a9fe34e08090a7300a40cd 100644
Binary files a/reconstruction_im/charis_VAE_bs_64.png and b/reconstruction_im/charis_VAE_bs_64.png differ
diff --git a/reconstruction_im/charis_VAE_bs_64_conv_64_64_128_128.png b/reconstruction_im/charis_VAE_bs_64_conv_64_64_128_128.png
index 928ef1498f0ad1286cc3d9139aebc553f7139607..9d7f6bd1c14a7d5c69ea18a12b4cc7babf83daf9 100644
Binary files a/reconstruction_im/charis_VAE_bs_64_conv_64_64_128_128.png and b/reconstruction_im/charis_VAE_bs_64_conv_64_64_128_128.png differ
diff --git a/reconstruction_im/charis_VAE_bs_64_ls_10_lr_1e_3.png b/reconstruction_im/charis_VAE_bs_64_ls_10_lr_1e_3.png
index a594af5ec3aeb327659ddcea73e6afa167ecbb41..46cba6316a85857ae39df116dfe7c14164458094 100644
Binary files a/reconstruction_im/charis_VAE_bs_64_ls_10_lr_1e_3.png and b/reconstruction_im/charis_VAE_bs_64_ls_10_lr_1e_3.png differ
diff --git a/reconstruction_im/charis_VAE_bs_64_ls_10_lr_5e_4.png b/reconstruction_im/charis_VAE_bs_64_ls_10_lr_5e_4.png
index c61be0bfd10db4861e278648b6f5ebed05096bf7..48b3175c0710c217c4d7717aa1fa958540218e55 100644
Binary files a/reconstruction_im/charis_VAE_bs_64_ls_10_lr_5e_4.png and b/reconstruction_im/charis_VAE_bs_64_ls_10_lr_5e_4.png differ
diff --git a/reconstruction_im/charis_VAE_bs_64_ls_15.png b/reconstruction_im/charis_VAE_bs_64_ls_15.png
index 5f33c6b9c5836971057647b29e3decce09e5ddc1..d40ff284bdadf54fb406f8e63e9da9941c5588e5 100644
Binary files a/reconstruction_im/charis_VAE_bs_64_ls_15.png and b/reconstruction_im/charis_VAE_bs_64_ls_15.png differ
diff --git a/reconstruction_im/charis_VAE_bs_64_ls_20.png b/reconstruction_im/charis_VAE_bs_64_ls_20.png
index aa5e9b6836ff7fbdfecd3cec43ab146b9b57e923..1cfc6aabf92566f8d0b14241c407b097416c2c8c 100644
Binary files a/reconstruction_im/charis_VAE_bs_64_ls_20.png and b/reconstruction_im/charis_VAE_bs_64_ls_20.png differ
diff --git a/reconstruction_im/charis_VAE_bs_64_ls_30.png b/reconstruction_im/charis_VAE_bs_64_ls_30.png
index c3cc443474dce18e6f39000582bb21b336f4145b..72ca3ac8195465b39539f6ea730ed12fb83e0808 100644
Binary files a/reconstruction_im/charis_VAE_bs_64_ls_30.png and b/reconstruction_im/charis_VAE_bs_64_ls_30.png differ
diff --git a/reconstruction_im/charis_VAE_bs_64_ls_40.png b/reconstruction_im/charis_VAE_bs_64_ls_40.png
index 1670cc17ff041e80530234a2d16550126553ed2b..0d32882024379b8e7101a2b9ebf55a248fcf36c5 100644
Binary files a/reconstruction_im/charis_VAE_bs_64_ls_40.png and b/reconstruction_im/charis_VAE_bs_64_ls_40.png differ
diff --git a/reconstruction_im/charis_VAE_bs_64_ls_5.png b/reconstruction_im/charis_VAE_bs_64_ls_5.png
index fc72e48d732b55f48489ce3f98b98fb13bac37e2..ac9a3c78a653a62075d592d874424c6b4dd04be3 100644
Binary files a/reconstruction_im/charis_VAE_bs_64_ls_5.png and b/reconstruction_im/charis_VAE_bs_64_ls_5.png differ
diff --git a/reconstruction_im/charis_VAE_bs_64_ls_50.png b/reconstruction_im/charis_VAE_bs_64_ls_50.png
index 8df74064debbace997044af694d3c3bf9e8db5f8..afd385eabda4b5ed37618a1e2d0f95b1b15794ab 100644
Binary files a/reconstruction_im/charis_VAE_bs_64_ls_50.png and b/reconstruction_im/charis_VAE_bs_64_ls_50.png differ
diff --git a/reconstruction_im/charis_beta_VAE_bs_256.png b/reconstruction_im/charis_beta_VAE_bs_256.png
index 16c26edfb3756b6b93865b38a9d70cfdd9feadf1..8cf93dfae0e85ed5200a6d9c176741d5fa1f768d 100644
Binary files a/reconstruction_im/charis_beta_VAE_bs_256.png and b/reconstruction_im/charis_beta_VAE_bs_256.png differ
diff --git a/reconstruction_im/charis_beta_VAE_bs_64.png b/reconstruction_im/charis_beta_VAE_bs_64.png
index 1be82dbc9602fc54f3597e1cb1a41a434b463543..eb6e0476c976e1ba5754238c2567328d5fa7783a 100644
Binary files a/reconstruction_im/charis_beta_VAE_bs_64.png and b/reconstruction_im/charis_beta_VAE_bs_64.png differ
diff --git a/reconstruction_im/charis_beta_VAE_bs_64_ls_15.png b/reconstruction_im/charis_beta_VAE_bs_64_ls_15.png
index dc336d17e847bd06ff3457de07184e0033ed146c..f6f00f249653d4b5801710a66fa56bee71302cdc 100644
Binary files a/reconstruction_im/charis_beta_VAE_bs_64_ls_15.png and b/reconstruction_im/charis_beta_VAE_bs_64_ls_15.png differ
diff --git a/reconstruction_im/charis_beta_VAE_bs_64_ls_20.png b/reconstruction_im/charis_beta_VAE_bs_64_ls_20.png
index 9291127082ea077774fa8db336c881fa98d4ac26..16e3a7bca0270c17c99d07cab0be73216df348e3 100644
Binary files a/reconstruction_im/charis_beta_VAE_bs_64_ls_20.png and b/reconstruction_im/charis_beta_VAE_bs_64_ls_20.png differ
diff --git a/reconstruction_im/charis_beta_VAE_bs_64_ls_5.png b/reconstruction_im/charis_beta_VAE_bs_64_ls_5.png
index 2f8022ded1ab3364ac3c8bfc356d29045aaf6706..c38f68dbd9fdb257fe8329175ac6b89ea9a203fd 100644
Binary files a/reconstruction_im/charis_beta_VAE_bs_64_ls_5.png and b/reconstruction_im/charis_beta_VAE_bs_64_ls_5.png differ
diff --git a/utils/training.py b/utils/training.py
index eac4cc9ed379105b7c2f74eef162f30b78133055..924aa53242b9bf35c52411a578c15119052444f5 100644
--- a/utils/training.py
+++ b/utils/training.py
@@ -63,6 +63,8 @@ class Trainer:
                        'kl_loss': []}
         self.global_iter = 0
         self.mean_epoch_loss = []
+        self.reconstruction_loss = []
+        self.kl_loss = []
         self.save_step = save_step
         self.expe_name = expe_name
         self.ckpt_dir = ckpt_dir
@@ -220,6 +222,7 @@ class Trainer:
 
         # with mse loss:
         recon_loss = F.mse_loss(recon_data, data, size_average=False).div(self.batch_size)
+        self.reconstruction_loss.append(recon_loss)
 
         prediction_loss = 0
         prediction_random_continue_loss = 0
@@ -267,6 +270,7 @@ class Trainer:
 
         # Calculate total kl value to record it
         kl_loss = kl_cont_loss + kl_disc_loss
+        self.kl_loss.append(kl_loss)
 
         # Calculate total loss
         if self.is_beta:
@@ -371,6 +375,8 @@ class Trainer:
         model_states = {'model': self.model.state_dict(), }
         optim_states = {'optim': self.optimizer.state_dict(), }
         states = {'loss': self.mean_epoch_loss,
+                  'recons_loss': self.reconstruction_loss,
+                  'kl_loss': self.kl_loss,
                   'iter': self.global_iter,
                   'model_states': model_states,
                   'optim_states': optim_states}
@@ -385,8 +391,10 @@ class Trainer:
         file_path = os.path.join(self.ckpt_dir, filename)
         if os.path.isfile(file_path):
             checkpoint = torch.load(file_path)
+            self.mean_epoch_loss = checkpoint['loss']
             if self.load_model_checkpoint:
-                self.mean_epoch_loss = checkpoint['loss']
+                self.reconstruction_loss = checkpoint['recons_loss']
+                self.kl_loss = checkpoint['kl_loss']
             self.global_iter = checkpoint['iter']
             self.model.load_state_dict(checkpoint['model_states']['model'])
             self.optimizer.load_state_dict(checkpoint['optim_states']['optim'])