[update]

export.py

@@ -555,7 +555,7 @@ def run(
         y = model(im)  # dry runs
     if half and not coreml:
         im, model = im.half(), model.half()  # to FP16
-    if os.environ['EXPORT_ARCH'] == 'ax620':
+    if os.environ.get("EXPORT_ARCH", "local") == 'ax620':
         metadata = {'stride': int(max(model.stride)), 'names': model.names}  # model metadata
     else:
         shape = tuple((y[0] if isinstance(y, tuple) else y).shape)  # model output shape

models/yolo.py

@@ -57,7 +57,7 @@ class Detect(nn.Module):
         z = []  # inference output
         for i in range(self.nl):
             x[i] = self.m[i](x[i])  # conv
-            if 'EXPORT_ARCH' in os.environ and os.environ['EXPORT_ARCH'] == 'ax620':
+            if os.environ.get("EXPORT_ARCH", "local") == 'ax620':
                 return x
             bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
             x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()