diff --git a/pts/trainer.py b/pts/trainer.py
index 0c95795..00728d8 100644
--- a/pts/trainer.py
+++ b/pts/trainer.py
@@ -65,13 +65,13 @@ class Trainer:
                     )
                     n_iter = epoch_no*self.num_batches_per_epoch + batch_no
                     writer.add_scalar('Loss/train', loss.item(), n_iter)
-                    for name, param in net.named_parameters():
-                        writer.add_histogram(name, param.clone().cpu().data.numpy(), n_iter)
 
                     loss.backward()
                     optimizer.step()
 
                     if self.num_batches_per_epoch == batch_no:
+                        for name, param in net.named_parameters():
+                            writer.add_histogram(name, param.clone().cpu().data.numpy(), n_iter)
                         break
 
             # mark epoch end time and log time cost of current epoch