diff --git a/hpvm/test/dnn_benchmarks/pytorch/dnn/_container.py b/hpvm/test/dnn_benchmarks/pytorch/dnn/_container.py
index 6f38094d1542f9e9a2bbc95a4f466db697c0ef6b..5918d960a745e5d245410acfac7c827b5b011f14 100644
--- a/hpvm/test/dnn_benchmarks/pytorch/dnn/_container.py
+++ b/hpvm/test/dnn_benchmarks/pytorch/dnn/_container.py
@@ -25,7 +25,7 @@ def make_conv_pool_activ(
 
 class Classifier(Module):
     def __init__(
-        self, convs: Sequential, linears: Sequential, use_softmax: bool = False
+        self, convs: Sequential, linears: Sequential, use_softmax: bool = True
     ):
         super().__init__()
         self.convs = convs