Solved by rmokady in post #2 
              
              
                I found some workaround. declare net.train() then for any Dropout and BatchNorm layers: 
    net.dropout.p = 0
    net.batchNorm.training = False

It seems to work.