Get stuck in fashion_mnist.py on pytorch 0.4.0

[cxmscb]

``
python fashion_mnist.py
waitGPU: Waiting for the following conditions, checking every 10 seconds.

• utilization <= 20
  waitGPU: Setting GPU to: [0]
  Traceback (most recent call last):
  File "fashion_mnist.py", line 63, in
  args.alpha_grad, args.scatter_grad, l1_proj=args.l1_proj)
  File "/home/songcb/convex_adversarial-master/examples/trainer.py", line 35, in train_robust
  out = model(Variable(X))
  File "/home/songcb/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 491, in call
  result = self.forward(*input, **kwargs)
  File "/home/songcb/.local/lib/python3.6/site-packages/torch/nn/modules/container.py", line 91, in forward
  input = module(input)
  File "/home/songcb/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 491, in call
  result = self.forward(*input, **kwargs)
  File "/home/songcb/.local/lib/python3.6/site-packages/torch/nn/modules/linear.py", line 55, in forward
  return F.linear(input, self.weight, self.bias)
  File "/home/songcb/.local/lib/python3.6/site-packages/torch/nn/functional.py", line 992, in linear
  return torch.addmm(bias, input, weight.t())
  RuntimeError: size mismatch, m1: [50 x 800], m2: [1568 x 100] at /pytorch/aten/src/THC/generic/THCTensorMathBlas.cu:249