diff --git a/yolov8/yolov8_cls_trt.py b/yolov8/yolov8_cls_trt.py
index 98c0846..ad75709 100644
--- a/yolov8/yolov8_cls_trt.py
+++ b/yolov8/yolov8_cls_trt.py
@@ -197,8 +197,10 @@ class YoLov8TRT(object):
 	    
 	    # Convert the image to row-major order, also known as "C order"
 	    image = np.ascontiguousarray(image)
+
+	    batch_data = np.expand_dims(image, axis=0)
 	    
-	    return image
+	    return batch_data
 
     def postprocess_cls(self, output_data):
         classes_ls = []