ONNX export self.training=False

This commit is contained in:
Glenn Jocher 2020-04-08 10:25:52 -07:00
parent 97780cfdb4
commit b6959a2f54
1 changed files with 1 additions and 0 deletions

View File

@ -143,6 +143,7 @@ class YOLOLayer(nn.Module):
self.anchor_wh = self.anchor_vec.view(1, self.na, 1, 1, 2) self.anchor_wh = self.anchor_vec.view(1, self.na, 1, 1, 2)
if ONNX_EXPORT: if ONNX_EXPORT:
self.training = False
self.create_grids((img_size[1] // stride, img_size[0] // stride)) # number x, y grid points 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'): def create_grids(self, ng=(13, 13), device='cpu'):