vame.train_model() error- UnboundLocalError

[samcentanni]

I'm stuck on the training part and I can't figure out what the source of this error is/ Here is the output for vame.train_model(config):

In [6]: vame.train_model(config)
Train Variational Autoencoder - model name: VAME

Using CUDA
GPU active: True
GPU used: Quadro RTX 6000
Latent Dimensions: 30, Time window: 30, Batch Size: 4096, Beta: 1, lr: 0.0010

Initialize train data. Datapoints 3618
Initialize test data. Datapoints 401
Scheduler step size: 3, Scheduler gamma: 0.50

Start training...
Epoch: 1

UnboundLocalError Traceback (most recent call last)
in
----> 1 vame.train_model(config)

~/anaconda3/envs/VAME/lib/python3.7/site-packages/vame-1.0-py3.7.egg/vame/model/rnn_vae.py in train_model(config)
338 FUTURE_STEPS, scheduler, MSE_REC_REDUCTION,
339 MSE_PRED_REDUCTION, KMEANS_LOSS, KMEANS_LAMBDA,
--> 340 TRAIN_BATCH_SIZE, noise)
341
342 current_loss, test_loss, test_list = test(test_loader, epoch, model, optimizer,

~/anaconda3/envs/VAME/lib/python3.7/site-packages/vame-1.0-py3.7.egg/vame/model/rnn_vae.py in train(train_loader, epoch, model, optimizer, anneal_function, BETA, kl_start, annealtime, seq_len, future_decoder, future_steps, scheduler, mse_red, mse_pred, kloss, klmbda, bsize, noise)
156
157 if future_decoder:
--> 158 print('Train loss: {:.3f}, MSE-Loss: {:.3f}, MSE-Future-Loss {:.3f}, KL-Loss: {:.3f}, Kmeans-Loss: {:.3f}, weight: {:.2f}'.format(train_loss / idx,
159 mse_loss /idx, fut_loss/idx, BETAkl_weightkullback_loss/idx, kl_weight*kmeans_losses/idx, kl_weight))
160 else:

UnboundLocalError: local variable 'idx' referenced before assignment

---

I traced this error and idx is defined in this model prior to this code. What am I missing?