I removed the transforms.ToPILImage() from train_dataset and the error is gone but now im getting this error:
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-12-b1d4464b941e> in <module>
      1 n_epochs = 10
      2 for epoch in range(n_epochs):
----> 3     train(epoch)

<ipython-input-9-df779ce045bf> in train(epoch)
      9         optimizer.zero_grad()
     10 
---> 11         output = model(data)
     12 
     13         loss = criterion(output, target)

~/.local/lib/python3.6/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    487             result = self._slow_forward(*input, **kwargs)
    488         else:
--> 489             result = self.forward(*input, **kwargs)
    490         for hook in self._forward_hooks.values():
    491             hook_result = hook(self, input, result)

<ipython-input-2-46779b84067d> in forward(self, x)
     32 
     33     def forward(self, x):
---> 34         x = self.features(x)
     35         x = x.view(x.size(0), -1)
     36         x = self.classifier(x)

~/.local/lib/python3.6/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    487             result = self._slow_forward(*input, **kwargs)
    488         else:
--> 489             result = self.forward(*input, **kwargs)
    490         for hook in self._forward_hooks.values():
    491             hook_result = hook(self, input, result)

~/.local/lib/python3.6/site-packages/torch/nn/modules/container.py in forward(self, input)
     90     def forward(self, input):
     91         for module in self._modules.values():
---> 92             input = module(input)
     93         return input
     94 

~/.local/lib/python3.6/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    487             result = self._slow_forward(*input, **kwargs)
    488         else:
--> 489             result = self.forward(*input, **kwargs)
    490         for hook in self._forward_hooks.values():
    491             hook_result = hook(self, input, result)

~/.local/lib/python3.6/site-packages/torch/nn/modules/conv.py in forward(self, input)
    318     def forward(self, input):
    319         return F.conv2d(input, self.weight, self.bias, self.stride,
--> 320                         self.padding, self.dilation, self.groups)
    321 
    322 

RuntimeError: Given groups=1, weight of size [32, 1, 5, 5], expected input[64, 3, 33, 33] to have 1 channels, but got 3 channels instead