ONNX export self.training=False

2020-04-08 10:25:52 -07:00 · 2020-04-08 10:25:52 -07:00 · b6959a2f54
parent 97780cfdb4
commit b6959a2f54
1 changed files with 1 additions and 0 deletions
--- a/models.py
+++ b/models.py
@ -143,6 +143,7 @@ class YOLOLayer(nn.Module):
        self.anchor_wh = self.anchor_vec.view(1, self.na, 1, 1, 2)

        if ONNX_EXPORT:
+            self.training = False
            self.create_grids((img_size[1] // stride, img_size[0] // stride))  # number x, y grid points

    def create_grids(self, ng=(13, 13), device='cpu'):