Size missmatch in linear layer when training on Bedrooms

[MunzerDw]

Hello!

First of all, thank you very much for publishing your code. I have preprocessed the data and want to start training, just as shown in the README. When I start training on bedroom room types, I am getting the following error:

This is happening at the "fc_class" linear layer in the "BaseAutoregressiveTransformer" module. The other room types work fine, and I am able to train with them (there is an overfitting issue though).
Do you know why the error above is happening? What can I change in the configuration file (i.e. network config) to fix it?

Thank you in advance for your help.

Best,
Munzer