CrossEntropyLoss target size

@J_Johnson’s suggestion of flattening the input sounds reasonable so did you try to fix the issue with it?