sets model to eval mode (required due to batchnorm layer)

fxia22 · Aug 18, 2018 · 4df9208 · 4df9208
1 parent 75dd61c
commit 4df9208
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/train_classification.py b/train_classification.py
@@ -74,6 +74,7 @@
         points = points.transpose(2,1)
         points, target = points.cuda(), target.cuda()
         optimizer.zero_grad()
+        classifier = classifier.train()
         pred, _ = classifier(points)
         loss = F.nll_loss(pred, target)
         loss.backward()
@@ -88,6 +89,7 @@
             points, target = Variable(points), Variable(target[:,0])
             points = points.transpose(2,1)
             points, target = points.cuda(), target.cuda()
+            classifier = classifier.eval()
             pred, _ = classifier(points)
             loss = F.nll_loss(pred, target)
             pred_choice = pred.data.max(1)[1]

diff --git a/train_segmentation.py b/train_segmentation.py
@@ -71,6 +71,7 @@
         points = points.transpose(2,1) 
         points, target = points.cuda(), target.cuda()   
         optimizer.zero_grad()
+        classifier = classifier.train()
         pred, _ = classifier(points)
         pred = pred.view(-1, num_classes)
         target = target.view(-1,1)[:,0] - 1
@@ -87,7 +88,8 @@
             points, target = data
             points, target = Variable(points), Variable(target)
             points = points.transpose(2,1) 
-            points, target = points.cuda(), target.cuda()   
+            points, target = points.cuda(), target.cuda()
+            classifier = classifier.eval()
             pred, _ = classifier(points)
             pred = pred.view(-1, num_classes)
             target = target.view(-1,1)[:,0] - 1