error in dataset creation of decoyMNIST

[avani17101]

The dataset creation and loading for decoy has a bug (getting the error below)
Could it be looked into @laura-rieger @csinva ?

python train_mnist_decoy.py

Traceback (most recent call last):
  File "train_mnist_decoy.py", line 107, in <module>
    complete_dataset = utils.TensorDataset(train_x_tensor,train_y_tensor) # create your datset
  File "/home/avani.gupta/miniconda3/envs/py3.8/lib/python3.8/site-packages/torch/utils/data/dataset.py", line 365, in __init__
    assert all(tensors[0].size(0) == tensor.size(0) for tensor in tensors), "Size mismatch between tensors"
AssertionError: Size mismatch between tensors