Hey I've been playing with the classes you define in your notebooks, and it seems to me that in the affine transformations you're multiplying weights and z in the wrong order. I changed it in my implementation to: (self.weights @ z.unsqueeze(-1)).squeeze(-1) , and made the appropriate change in the _inverse method too
Hey I've been playing with the classes you define in your notebooks, and it seems to me that in the affine transformations you're multiplying
weightsandzin the wrong order. I changed it in my implementation to:(self.weights @ z.unsqueeze(-1)).squeeze(-1), and made the appropriate change in the_inversemethod too