Switch net to train mode in train.py

Phaeton-531 · Sep 26, 2018 · ab255c4 · ab255c4
1 parent 3ab8272
commit ab255c4
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/train.py b/train.py
@@ -54,6 +54,7 @@ def train_net(net,
 
     for epoch in range(epochs):
         print('Starting epoch {}/{}.'.format(epoch + 1, epochs))
+        net.train()
 
         # reset the generators
         train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_scale)