diff --git a/flaxdiff/data/online_loader.py b/flaxdiff/data/online_loader.py
index 63ff0ab..74ad0c2 100644
--- a/flaxdiff/data/online_loader.py
+++ b/flaxdiff/data/online_loader.py
@@ -111,8 +111,8 @@ def map_sample(
             "original_width": original_width,
         })
     except Exception as e:
-        print(f"Error maping sample {url}", e)
-        traceback.print_exc()
+        # print(f"Error maping sample {url}", e)
+        # traceback.print_exc()
         # error_queue.put_nowait({
         #     "url": url,
         #     "caption": caption,
diff --git a/flaxdiff/trainer/simple_trainer.py b/flaxdiff/trainer/simple_trainer.py
index 594e06c..b8387f0 100644
--- a/flaxdiff/trainer/simple_trainer.py
+++ b/flaxdiff/trainer/simple_trainer.py
@@ -437,5 +437,6 @@ def train_loop(current_step, pbar: tqdm.tqdm, train_state, rng_state):
                         "train/epoch": current_epoch,
                     }, step=current_step)
                 print(colored(f"\n\tEpoch {current_epoch} completed. Avg Loss: {avg_loss}, Time: {total_time:.2f}s, Best Loss: {self.best_loss}", 'green'))
+        print("Training done")
         self.save(epochs)
         return self.state
diff --git a/setup.py b/setup.py
index 28c047e..30c0227 100644
--- a/setup.py
+++ b/setup.py
@@ -11,7 +11,7 @@
 setup(
     name='flaxdiff',
     packages=find_packages(),
-    version='0.1.35.3',
+    version='0.1.35.4',
     description='A versatile and easy to understand Diffusion library',
     long_description=open('README.md').read(),
     long_description_content_type='text/markdown',