Bug when training DIVE on CPU

[RukuangHuang]

When GPU is not available, training DIVE (or possibly other models with more than one KL term), the ALREADY_EXISTS error will be raised after the first epoch.

Epoch 1/2
189/189 [==============================] - 1141s 6s/step - loss: 22154.2715 - group_covs_kernel: 0.1143 - covs_dev_decoder: 0.0030 - ll_loss: 22154.1484 - kl_loss: 0.0000e+00 - lr: 0.0050
Epoch 2/2
2024-01-26 12:04:48.265028: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at variable_ops.cc:118 : ALREADY_EXISTS: Resource __per_step_0/DIVE/kl_loss/AddN/tmp_var/frame:0/iter:0/N10tensorflow19TemporaryVariableOp6TmpVarE
2024-01-26 12:04:48 ERROR osl-dynamics [pipeline.py:131:run_pipeline]: Graph execution error:

Resource __per_step_0/DIVE/kl_loss/AddN/tmp_var/frame:0/iter:0/N10tensorflow19TemporaryVariableOp6TmpVarE
         [[{{node DIVE/kl_loss/AddN/tmp_var}}]] [Op:__inference_train_function_113708]

[RukuangHuang]

Update: This only happens when use_tfrecord = True. So the temporary solution is to avoid using TFRecords when GPU is not available.