diff --git a/examples/classifier_compression/compress_classifier.py b/examples/classifier_compression/compress_classifier.py index 5ba8222866a1a8cbbaad480c39b90ade398aac24..d0b53db6512d6ebda4043a9789342ebf185f725b 100755 --- a/examples/classifier_compression/compress_classifier.py +++ b/examples/classifier_compression/compress_classifier.py @@ -509,6 +509,8 @@ def train(train_loader, model, criterion, optimizer, epoch, if lc.name not in losses: losses[lc.name] = tnt.AverageValueMeter() losses[lc.name].add(lc.value.item()) + else: + losses[OVERALL_LOSS_KEY].add(loss.item()) # Compute the gradient and do SGD step optimizer.zero_grad()