Running Hardware Optimized NVIDIA Optical Flow using OpenCV Python

Using the code from the following link https://devblogs.nvidia.com/opencv-optical-flow-algorithms-with-nvidia-turing-gpus/ I tried to estimate the optical flow of a sample video using the below code.

import numpy as np 
import cv2

def testing(vid_path):
    cap = cv.VideoCapture(vid_path)
    i = 0
    frames = []
    flow = []
    size = None
    while(cap.isOpened()):
        success, frame = cap.read()
        height, width, layers = frame.shape
        size = (width,height)
        if success:
            frames.append(frame)
            if i == 0:
                i += 1
                continue
            else:
                print(frames[i-1].shape, frames[i].shape)
                nvof = cv.cuda_NvidiaOpticalFlow_1_0.create(frame.shape[1], frame.shape[0], 5, False, False, False, 0)
                flow = nvof.calc(frames[i-1], frames[i], None)
                flowUpSampled = nvof.upSampler(flow[0], frame.shape[1], frame.shape[0], nvof.getGridSize(), None)
                flow.append(flowUpSampled)
                i += 1
                nvof.collectGarbage()
    cap.release()
    out = cv.VideoWriter('video.avi',cv.VideoWriter_fourcc(*'DIVX'), i, size)
    for i in range(len(flow)):
        out.write(flow[i])
    out.release()

The function testing has all the functionality and takes in the path of the video for which Optical Flow had to be estimated. I want to run it on my GPU which is an RTX 2070. The above code gives out the following error -

 (240, 320, 3) (240, 320, 3)
---------------------------------------------------------------------------
error                                     Traceback (most recent call last)
<ipython-input-14-13554636257c> in <module>
----> 1 testing('C:/Users/.../Downloads/TestingVid.mpg')

<ipython-input-13-207420f4d52a> in testing(vid_path)
 17                 print(frames[i-1].shape, frames[i].shape)
 18                 nvof = cv.cuda_NvidiaOpticalFlow_1_0.create(frame.shape[1], frame.shape[0], 5, False, False, False, 0)
---> 19                 flow = nvof.calc(frames[i-1], frames[i], None)
 20                 flowUpSampled = nvof.upSampler(flow[0], frame.shape[1], frame.shape[0], nvof.getGridSize(), None)
 21                 flow.append(flowUpSampled)

error: OpenCV(4.3.0-dev) C:/Users/.../Downloads/opencv/opencv-master/modules/core/src/cuda/gpu_mat.cu:224: error: (-217:Gpu API call) invalid pitch argument in function 'cv::cuda::GpuMat::upload'

I don't understand the error properly. I tried to look at the source code, but I couldn't find what was causing that error. CUDA Toolkit version 10.2 Python 3.8 TIA.

Edit - Understood the error. Found it. My GPU isn't large enough to have a pitch value for a (240, 320, 3) frame in the GPU memory, hence that error. Once I change it to grayscale, the error vanishes.