jupitor notebookでの不具合について

前提・実現したいこと

ここに質問の内容を詳しく書いてください。
google colabでは問題なく動作しておりましたが、同一のコードでvisual studio Codeのjupitor notebookで動作させると以下のエラーメッセージが出力されてしまいます。理由及び対処が分からなくアドバイスを頂けますと幸いでございます。

発生している問題・エラーメッセージ

エラーメッセージ
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-7-7daad3dd4759> in <module>()
     15 
     16 # Get a batch of training data
---> 17 inputs, classes = next(iter(dataloaders['train']))
     18 
     19 # Make a grid from batch

/home/imagawa/.local/lib/python3.6/site-packages/torch/utils/data/dataloader.py in __next__(self)
    519             if self._sampler_iter is None:
    520                 self._reset()
--> 521             data = self._next_data()
    522             self._num_yielded += 1
    523             if self._dataset_kind == _DatasetKind.Iterable and \

/home/imagawa/.local/lib/python3.6/site-packages/torch/utils/data/dataloader.py in _next_data(self)
    559     def _next_data(self):
    560         index = self._next_index()  # may raise StopIteration
--> 561         data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
    562         if self._pin_memory:
    563             data = _utils.pin_memory.pin_memory(data)

/home/imagawa/.local/lib/python3.6/site-packages/torch/utils/data/_utils/fetch.py in fetch(self, possibly_batched_index)
     42     def fetch(self, possibly_batched_index):
     43         if self.auto_collation:
---> 44             data = [self.dataset[idx] for idx in possibly_batched_index]
     45         else:
     46             data = self.dataset[possibly_batched_index]

/home/imagawa/.local/lib/python3.6/site-packages/torch/utils/data/_utils/fetch.py in <listcomp>(.0)
     42     def fetch(self, possibly_batched_index):
     43         if self.auto_collation:
---> 44             data = [self.dataset[idx] for idx in possibly_batched_index]
     45         else:
     46             data = self.dataset[possibly_batched_index]

/home/imagawa/.local/lib/python3.6/site-packages/torch/utils/data/dataset.py in __getitem__(self, idx)
    309 
    310     def __getitem__(self, idx):
--> 311         return self.dataset[self.indices[idx]]
    312 
    313     def __len__(self):

/home/imagawa/.local/lib/python3.6/site-packages/torch/utils/data/dataset.py in __getitem__(self, idx)
    309 
    310     def __getitem__(self, idx):
--> 311         return self.dataset[self.indices[idx]]
    312 
    313     def __len__(self):

/home/imagawa/.local/lib/python3.6/site-packages/torchvision/datasets/folder.py in __getitem__(self, index)
    232         sample = self.loader(path)
    233         if self.transform is not None:
--> 234             sample = self.transform(sample)
    235         if self.target_transform is not None:
    236             target = self.target_transform(target)

/home/imagawa/.local/lib/python3.6/site-packages/torchvision/transforms/transforms.py in __call__(self, img)
     58     def __call__(self, img):
     59         for t in self.transforms:
---> 60             img = t(img)
     61         return img
     62 

/home/imagawa/.local/lib/python3.6/site-packages/torchvision/transforms/transforms.py in __call__(self, pic)
     95             Tensor: Converted image.
     96         """
---> 97         return F.to_tensor(pic)
     98 
     99     def __repr__(self):

/home/imagawa/.local/lib/python3.6/site-packages/torchvision/transforms/functional.py in to_tensor(pic)
    127     mode_to_nptype = {'I': np.int32, 'I;16': np.int16, 'F': np.float32}
    128     img = torch.from_numpy(
--> 129         np.array(pic, mode_to_nptype.get(pic.mode, np.uint8), copy=True)
    130     )
    131 

TypeError: __array__() takes 1 positional argument but 2 were given

### 該当のソースコード
def imshow(inp, title=None):
    """Imshow for Tensor."""
    print(type(inp))
    inp = inp.numpy().transpose((1, 2, 0))
    print(type(inp))
    mean= np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    inp = std * inp + mean
    inp = np.clip(inp, 0, 1)
    plt.imshow(inp)
    if title is not None:
        plt.title(title)
    plt.pause(0.001)  # pause a bit so that plots are updated


# Get a batch of training data
inputs, classes = next(iter(dataloaders['train']))

# Make a grid from batch
out = torchvision.utils.make_grid(inputs)

imshow(out, title=[class_names[x] for x in classes])