We currently quick hacked the reshaping into the training script after keras changes how it expects the input data