Loadweights.py

# First, we need the same structure:
new_model = tl.Serial(
    tl.Embedding(data.vocab_size(vocab_file='en_8k.subword'), d_feature=256),
    tl.Mean(axis=1),
    tl.Dense(2),
    tl.LogSoftmax()
)
# Then, we load the weights:
new_model.init_from_file(file_name="/root/output_dir/model.pkl.gz", weights_only=True) # Only load weights
# Same result as before (I used a helper function for simplicity)
print("The sentiment is: ", parse_sentiment("Very bad movie", new_model))