Hi,
I haven’t gone through all the code but have you tried either removing the dropout during training or keeping the dropout during testing? This seem to be the only module affected by model.eval() in your model.