Error on normalize_target method of LadderVAE class

[mese79]

I was running the EvaluationOnTestImage notebook, and I've got an error regarding the two tensor are not on the same device.
It happens in normilize_target method: code
As a quick fix, I did this:

def normalize_target(self, target, batch=None):
    device = target.device
    return (target - self.data_mean.to(device)) / self.data_std.to(device)

I'm not sure if this is a bug or something's wrong with the notebook that I've run.