Solved by Alexey_Demyanchuk in post #2 
              
              
                Hi, I am not sure, but probably here you don’t need to put batch on device twice. 

Also at this line your model (autoencod) is on device, but data is not