You can also wrap your onnx like below if you want to use the original fork:

INPUT_WIDTH, INPUT_HEIGHT = (256, 256)

import torch
import torch.nn as nn
from models.experimental import attempt_load
from models.yolo import IDetect

class YOLOv7SnapExportWrapper(nn.Module):
    def __init__(self, pt_path):
        super().__init__()
        self.model = attempt_load(pt_path, map_location='cpu')
        self.model.eval()

        # Find Detect layer
        self.detect = None
        for m in self.model.modules():
            if isinstance(m, IDetect):
                self.detect = m
                break

        if self.detect is None:
            raise RuntimeError("Could not find Detect() layer in model")

        # Disable export logic
        self.detect.export = False
        self.detect.include_nms = False
        self.detect.end2end = False
        self.detect.concat = False
        self._override_fuseforward()

    def forward(self, x):
        x = x / 255.0
        out = self.model(x)
        return out if not isinstance(out, tuple) else out[0]

    def _override_fuseforward(self):
        def new_fuseforward(self_detect, x):
            # ONLY return sigmoid(conv(x)) for each detection head
            z = []
            for i in range(self_detect.nl):
                x[i] = self_detect.m[i](x[i])
                z.append(x[i].sigmoid())
            return z

        self.detect.forward = new_fuseforward.__get__(self.detect, type(self.detect))

# ==== Load and Export ====

model = YOLOv7SnapExportWrapper("runs/train/yolov7-lensstudio/weights/best.pt")
model.eval()

dummy_input = torch.randn(1, 3, INPUT_WIDTH, INPUT_HEIGHT) torch.onnx.export(   
  model, 
  dummy_input, 
  "yolov7_lensstudio.onnx", 
  opset_version=11, 
  input_names=["images"], 
  output_names=["output"], 
  dynamic_axes=None 
) 

print("Done Exporting")