Fix batch-norm layers in PyTorchYolo._get_losses

Signed-off-by: Beat Buesser <[email protected]>

Author: Beat Buesser

art/estimators/object_detection/pytorch_yolo.py

@@ -253,6 +253,8 @@ def _get_losses(
         import torch  # lgtm [py/repeated-import]
 
         self._model.train()
+        self.set_batchnorm(train=False)
+        self.set_dropout(train=False)
 
         # Apply preprocessing
         if self.all_framework_preprocessing: