From 79fba1fc260c1c29f8227ffc18f99603ea1c6b35 Mon Sep 17 00:00:00 2001 From: "Dr. Kashif Rasul" Date: Thu, 16 Jan 2020 18:24:05 +0100 Subject: [PATCH] log weights at the end of each epoch --- pts/trainer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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