terminate called after throwing an instance of 'c10::Error'
  what():  input.ndimension() == 4 ASSERT FAILED at /pytorch/torch/csrc/api/src/nn/modules/conv.cpp:96, please report a bug to PyTorch. (forward at /pytorch/torch/csrc/api/src/nn/modules/conv.cpp:96)
frame #0: std::function<std::string ()>::operator()() const + 0x11 (0x7fdfa993ccd1 in /lib/libtorch/lib/libc10.so)
frame #1: c10::Error::Error(c10::SourceLocation, std::string const&) + 0x2a (0x7fdfa993c59a in /lib/libtorch/lib/libc10.so)
frame #2: torch::nn::Conv2dImpl::forward(at::Tensor) + 0x22c (0x7fdfb23109ec in /lib/libtorch/lib/libtorch.so.1)
frame #3: Net::forward(at::Tensor) + 0xdf (0x7fdfb2c8cebd in ./mnist)
frame #4: void train<torch::data::DataLoader<torch::data::datasets::MapDataset<torch::data::datasets::MapDataset<torch::data::datasets::MapDataset<torch::data::datasets::MNIST, Normalize>, torch::data::transforms::TensorLambda<at::Tensor> >, torch::data::transforms::Stack<torch::data::Example<at::Tensor, at::Tensor> > >, torch::data::samplers::RandomSampler> >(int, Options const&, Net&, c10::Device, torch::data::DataLoader<torch::data::datasets::MapDataset<torch::data::datasets::MapDataset<torch::data::datasets::MapDataset<torch::data::datasets::MNIST, Normalize>, torch::data::transforms::TensorLambda<at::Tensor> >, torch::data::transforms::Stack<torch::data::Example<at::Tensor, at::Tensor> > >, torch::data::samplers::RandomSampler>&, torch::optim::SGD&, unsigned long) + 0x15c (0x7fdfb2c91ebf in ./mnist)
frame #5: main + 0x512 (0x7fdfb2c8547f in ./mnist)
frame #6: __libc_start_main + 0xe7 (0x7fdfa8d81b97 in /lib/x86_64-linux-gnu/libc.so.6)
frame #7: _start + 0x2a (0x7fdfb2c83c6a in ./mnist)

Aborted (core dumped)