Convert tensorflow code to pytorch

2

The mapping should be:

  • tf.reduce_mean -> tensor.mean
  • tf.expand_dims -> tensor.expand
  • tf.transpose -> tensor.permute

Let us know, if you have any trouble.