Error when trying to train mnist poincare experiment

[saniazahan]

I got the following error while testing the mnist poincare experiment. If I train for 79 epoch the error doesn't occur. Could you please suggest what is the problem

[Ingmar-Bergman]

Try to replace 'input = input.unsqueeze(-2).expand(*input.shape[:-(len(input.shape) - 2)], self.out_features, self.in_features)'
by 'input = input.unsqueeze(-2).expand(*input.shape[:-1], self.out_features, self.in_features)'.

It worked for me.