TypeError: __init__() got an unexpected keyword argument 'new_step_api'

Hi
when i run the following Pytorch code in Jupyter notebook, i am getting the error as TypeError: init() got an unexpected keyword argument 'new_step_api'

Could you please help me on finding the solution for this problem ?

import gym
import math
import random
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from collections import namedtuple, deque
from itertools import count
from PIL import Image

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torchvision.transforms as T

env = gym.make('CartPole-v0', new_step_api=True, render_mode='single_rgb_array').unwrapped

set up matplotlib

is_ipython = 'inline' in matplotlib.get_backend()
if is_ipython:
from IPython import display

plt.ion()

if gpu is to be used

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

cc @vmoens @nairbv

Start by reproducing the Jupyter notebook snippet that calls gym.make for CartPole-v0 with new_step_api and render_mode. Check the tutorial's expected dependency versions and the installed Gym API; done means the example runs without this TypeError or clearly documents the compatible setup.