Thanks for your response.
I am using the following repo with customizations for my data:
  
      
      github.com
  
  
    

fabio-deep/Variational-Capsule-Routing 4


  Official Pytorch code for (AAAI 2020) paper "Capsule Routing via Variational Bayes", https://arxiv.org/pdf/1905.11455.pdf

  
  
    
    
  
  


When I run the main.py I get the same error as I get while compiling in console line by line…
Following is the complete stack of error:
Traceback (most recent call last):
File “F:/Richa/Richa_Python_codes/Variational-Capsule-Routing/src/main_r.py”, line 136, in 
score = main(parser.parse_known_args()[0])
File “F:/Richa/Richa_Python_codes/Variational-Capsule-Routing/src/main_r.py”, line 107, in main
score = train(model, dataloaders, args)
File “F:\Richa\Richa_Python_codes\Variational-Capsule-Routing\src\train.py”, line 103, in train
test_loss, test_acc = evaluate(model, args, dataloaders[‘test’])
File “F:\Richa\Richa_Python_codes\Variational-Capsule-Routing\src\evaluate.py”, line 13, in evaluate
for i, (inputs, labels) in enumerate(dataloader):
File “C:\Users\user.conda\envs\gpu-test4\lib\site-packages\torch\utils\data\dataloader.py”, line 435, in next
data = self._next_data()
File “C:\Users\user.conda\envs\gpu-test4\lib\site-packages\torch\utils\data\dataloader.py”, line 475, in _next_data
data = self.dataset_fetcher.fetch(index)  # may raise StopIteration
File “C:\Users\user.conda\envs\gpu-test4\lib\site-packages\torch\utils\data_utils\fetch.py”, line 44, in fetch
data = [self.dataset[idx] for idx in possibly_batched_index]
File “C:\Users\user.conda\envs\gpu-test4\lib\site-packages\torch\utils\data_utils\fetch.py”, line 44, in 
data = [self.dataset[idx] for idx in possibly_batched_index]
File “F:\Richa\Richa_Python_codes\Variational-Capsule-Routing\src\datasets_r.py”, line 120, in getitem
image = self.transform(self.data[idx])
File “C:\Users\user.conda\envs\gpu-test4\lib\site-packages\torchvision\transforms\transforms.py”, line 67, in call
img = t(img)
File “C:\Users\user.conda\envs\gpu-test4\lib\site-packages\torch\nn\modules\module.py”, line 727, in call_impl
result = self.forward(*input, **kwargs)
File “C:\Users\user.conda\envs\gpu-test4\lib\site-packages\torchvision\transforms\transforms.py”, line 226, in forward
return F.normalize(tensor, self.mean, self.std, self.inplace)
File “C:\Users\user.conda\envs\gpu-test4\lib\site-packages\torchvision\transforms\functional.py”, line 284, in normalize
tensor.sub(mean).div(std)
RuntimeError: The size of tensor a (3) must match the size of tensor b (2) at non-singleton dimension 0