Convert a code from tensorflow/keras to pytorch frame work

[sumayya]

I print the shape after use function (permute)