DIEN_TF2 .contrib is used by FP16 and mode=test

[eliommar]

Is there any way to solve the problem related to .contrib module not found using Tensorflow >=2?
I'm using --mode=test --data_type=FP16 --embedding_device=cpu.
Output:
Traceback (most recent call last):
File "script/train.py", line 413, in
test(model_type=args.model, seed=SEED, batch_size=args.batch_size, data_type=args.data_type)
File "script/train.py", line 358, in test
fp32_variables = [var_name for var_name, _ in tf.contrib.framework.list_variables(model_path)]
AttributeError: module 'tensorflow' has no attribute 'contrib'
I tried with t2_comapt_v1 and it doesnt solve the problem

[eliommar]

Steps to reproduce:
pip install tensorflow (this will install TF 2.4)
python script/train.py --mode=test --data_type=FP16 --batch_size=128 --model=DIEN --embedding_device=cpu