RuntimeError                              Traceback (most recent call last)
 in ()
6
7 for epoch in range(epochs):
----> 8     for inputs, labels in trainloader:
9         steps += 1
10         inputs, labels = inputs.to(device), labels.to(device)
~/.local/lib/python3.6/site-packages/torch/utils/data/dataloader.py in next(self)
613         if self.num_workers == 0:  # same-process loading
614             indices = next(self.sample_iter)  # may raise StopIteration
–> 615             batch = self.collate_fn([self.dataset[i] for i in indices])
616             if self.pin_memory:
617                 batch = pin_memory_batch(batch)
~/.local/lib/python3.6/site-packages/torch/utils/data/dataloader.py in default_collate(batch)
230     elif isinstance(batch[0], container_abcs.Sequence):
231         transposed = zip(*batch)
–> 232         return [default_collate(samples) for samples in transposed]
233
234     raise TypeError((error_msg.format(type(batch[0]))))
~/.local/lib/python3.6/site-packages/torch/utils/data/dataloader.py in (.0)
230     elif isinstance(batch[0], container_abcs.Sequence):
231         transposed = zip(*batch)
–> 232         return [default_collate(samples) for samples in transposed]
233
234     raise TypeError((error_msg.format(type(batch[0]))))
~/.local/lib/python3.6/site-packages/torch/utils/data/dataloader.py in default_collate(batch)
207             storage = batch[0].storage().new_shared(numel)
208             out = batch[0].new(storage)
–> 209         return torch.stack(batch, 0, out=out)
210     elif elem_type.module == ‘numpy’ and elem_type.name != 'str’ 
211             and elem_type.name != ‘string_’:
RuntimeError: invalid argument 0: Sizes of tensors must match except in dimension 0. Got 344 and 339 in dimension 3 at /pytorch/aten/src/TH/generic/THTensorMoreMath.cpp:1307