Hello,
Thank you very much for your reply.
Could you be more specific about your code? I am new to Python, and plus I am new to Pyro, so I am not sure how I will integrate your answer to my code.
What are the input, batch_size, and dim1 in your code?