~~Implement encoder/decoder constructor.~~ ~~Implement loglikelihoods etc. for 4D tensors.~~ ~~Make sure everything works on GPU.~~