Failure example.ipynb conversion of keras model to pytorch

[wfong443]

Probably related to issue #8. Running the code in the notebook example.ipynb showed differences between the keras and pytorch model.

The issue is that the Keras Flatten function does a permutation (see https://github.com/keras-team/keras/blob/master/keras/layers/core.py#L507) of the input data when the data_format is set to 'channel_first' which unnecessary since the data is provided in channel_first format. Essentially, Flatten permutes the data into the incorrect format which leads to the observed differences.

The fix is to specify the data format to prevent the permutation:
model.add(Flatten(data_format='channels_last'))
This creates an instance of Flatten that will not permute the data. See attached file.
example.ipynb.zip