TF2 changed the softmax cross entropy function, they removed the variable dim and recommended to use axis.
dim
axis
fmpr
