I met a strange problem in the PyTorch official Reinforcement Learning (DQN) tutorial  4.
It is OK when I run the reinforcement_q_learning.ipynb on my macbookpro(OSX), however, when I ran the last cell of it on my PC (Ubuntu16.04+1080ti+Cuda8.0+PyTorch+Gym), it errors out:
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-8-c9b5f0b71a6f> in <module>()
  3     # Initialize the environment and state
  4     env.reset()
----> 5     last_screen = get_screen()
  6     current_screen = get_screen()
  7     state = current_screen - last_screen

<ipython-input-5-3aa424675ed4> in get_screen()
 34     screen = torch.from_numpy(screen)
 35     # Resize, and add a batch dimension (BCHW)
---> 36     return resize(screen).unsqueeze(0).type(Tensor)
 37 
 38 env.reset()

/usr/local/lib/python2.7/dist-packages/torchvision/transforms.pyc in __call__(self, img)
 27     def __call__(self, img):
 28         for t in self.transforms:
---> 29             img = t(img)
 30         return img
 31 

/usr/local/lib/python2.7/dist-packages/torchvision/transforms.pyc in __call__(self, pic)
 77             pic = pic.mul(255).byte()
 78         if torch.is_tensor(pic):
---> 79             npimg = np.transpose(pic.numpy(), (1, 2, 0))
 80         assert isinstance(npimg, np.ndarray), 'pic should be Tensor or ndarray'
 81         if npimg.shape[2] == 1:

/home/diaosiki/.local/lib/python2.7/site-packages/numpy/core/fromnumeric.pyc in transpose(a, axes)
548 
549     """
--> 550     return _wrapfunc(a, 'transpose', axes)
551 
552 

/home/diaosiki/.local/lib/python2.7/site-packages/numpy/core/fromnumeric.pyc in _wrapfunc(obj, method, *args, **kwds)
 55 def _wrapfunc(obj, method, *args, **kwds):
 56     try:
---> 57         return getattr(obj, method)(*args, **kwds)
 58 
 59     # An AttributeError occurs if the object does not have

ValueError: axes don't match array

I could not figure out what’s the problem. Since the tutorial run on OSX, it should be my PC environment problem, yet other tutorials like Transfer Learning tutorial 1 run perfectly on my PC.
Thanks for any help!