Layer Normalization is not enabled in train.py, while admitted defaulty in model.py using TopK

[InuyashaYang]

In train.py, we observe that x is processed in FastAutoencoder.forward in this way:

x = x - pre_bias
latents_pre_act = F.linear(x, weight, latent_bias)

which is only zero-centered, not normalized

while in the model.py

# NOTE: hacky way to determine if normalization is enabled
normalize = activation_class_name == "TopK" 

this causes the model trained by the script to perform badly when using model.py to infer, I wonder this is a misuse or a part left to be implemented