Assertion Error in Jetson-Reinforcement

Hi, I am following the tutorial for Jetson Reinforcement Learning, https://github.com/dusty-nv/jetson-reinforcement. I am getting AssertionError in Jupyter Notebook. I am not sure what is causing the error. Please help.

I am using Pytorch 1.4.0, in Ubuntu 16.04.


AssertionError Traceback (most recent call last)
in ()
9 # Select and perform an action
10 action = select_action(state)
—> 11 _, reward, done, _ = env.step(action[0, 0])
12 #action = q_value.max(1)[1].data[0] -> action = q_value.max(1)[1].item()
13 #losses.append(loss.data[0]) -> losses.append(loss.item())

/usr/local/lib/python2.7/dist-packages/gym/core.pyc in step(self, action)
97 info (dict): contains auxiliary diagnostic information (helpful for debugging, and sometimes learning)
98 “”"
—> 99 return self._step(action)
100
101 def reset(self):

/usr/local/lib/python2.7/dist-packages/gym/envs/classic_control/cartpole.pyc in _step(self, action)
54
55 def _step(self, action):
—> 56 assert self.action_space.contains(action), “%r (%s) invalid”%(action, type(action))
57 state = self.state
58 x, x_dot, theta, theta_dot = state

AssertionError: tensor(1, device=‘cuda:0’) (<class ‘torch.Tensor’>) invalid