"How can I deploy a classification model outside of predict.py?

[cadpm]

Hi,

I tried deploy classification in my own program. But i have this error:

import torch
import cv2

model_yolo_class = torch.hub.load('yolov5', "custom", "yolov5s-cls.pt", source='local')  # load from PyTorch Hub
cap = cv2.VideoCapture(0)

while True:
    ret, frame = cap.read()
    detect_yolo_class = model_yolo_class(frame)

  File "C:\Users\dead_\Desktop\YoloforJetson\yolo.py", line 28, in <module>
    detect_yolo_class = model_yolo_class(frame)
  File "C:\Users\dead_\Desktop\YoloforJetson\lib\site-packages\torch\nn\modules\module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Users\dead_\Desktop\YoloforJetson\yolov5\models\common.py", line 520, in forward
    b, ch, h, w = im.shape  # batch, channel, height, width
ValueError: not enough values to unpack (expected 4, got 3)

I have been searching for information, but I haven't found anything anywhere. If there's any way to implement it, could you please share some example code?

I'm working with a Jetson Nano, and YOLOv8, with which I have a bit more experience, doesn't work on this device.
I wanted to implement my model on this device.